test database connection

Author: imabdulbasit

Cargo.lock

@@ -5,6 +5,9 @@ authors = ["Espresso Systems <[email protected]>"]
 edition = "2021"
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
+[dev-dependencies]
+hotshot-query-service = { git = "https://github.com/EspressoSystems/hotshot-query-service", tag = "rc-0.1.44" }
+
 [dependencies]
 anyhow = "1"
 async-compatibility-layer = { version = "1.1", default-features = false, features = [
@@ -21,7 +24,7 @@ hotshot-events-service = { git = "https://github.com/EspressoSystems/hotshot-eve
 hotshot-types = { git = "https://github.com/EspressoSystems/HotShot.git", tag = "rc-0.5.61" }
 serde = { version = "1.0", features = ["derive"] }
 snafu = "0.8"
-sqlx = { version = "0.7.4", features = [ "postgres"] }
+sqlx = { version = "0.7.4", features = [ "postgres", "macros" ] }
 surf-disco = "0.9"
 tide-disco = "0.9"
 toml = "0.8.14"

Cargo.toml

@@ -1,3 +1,4 @@
 CREATE TABLE test (
-    id     SERIAL PRIMARY KEY,
-);
+    id SERIAL PRIMARY KEY,
+    str text NOT NULL
+);

migration/test.sql

@@ -1,4 +1,10 @@
-use sqlx::{pool::PoolConnection, postgres::PgPoolOptions, Error, PgPool, Postgres};
+use anyhow::Context;
+use sqlx::{
+    pool::PoolConnection,
+    postgres::{PgConnectOptions, PgPoolOptions},
+    ConnectOptions, Error, PgPool, Postgres,
+};
+use tide_disco::Url;
 
 use crate::DatabaseOptions;
 
@@ -8,13 +14,41 @@ impl PostgresClient {
     pub async fn connect(opts: DatabaseOptions) -> anyhow::Result<Self> {
         let DatabaseOptions {
             url,
+            host,
+            port,
+            db_name,
+            username,
+            password,
             max_connections,
             acquire_timeout,
             migrations,
         } = opts;
 
         let mut options = PgPoolOptions::new();
 
+        let postgres_url: Url = match url {
+            Some(url) => url.parse()?,
+            None => {
+                let host = host.context("host not provided")?;
+                let port = port.context("port not provided")?;
+                let mut connect_opts = PgConnectOptions::new().host(&host).port(port);
+
+                if let Some(username) = username {
+                    connect_opts = connect_opts.username(&username);
+                }
+
+                if let Some(password) = password {
+                    connect_opts = connect_opts.password(&password);
+                }
+
+                if let Some(db_name) = db_name {
+                    connect_opts = connect_opts.database(&db_name);
+                }
+
+                connect_opts.to_url_lossy()
+            }
+        };
+
         if let Some(max_connections) = max_connections {
             options = options.max_connections(max_connections);
         }
@@ -23,7 +57,7 @@ impl PostgresClient {
             options = options.acquire_timeout(acquire_timeout);
         }
 
-        let connection = options.connect(&url).await?;
+        let connection = options.connect(postgres_url.as_str()).await?;
 
         if migrations {
             sqlx::migrate!("./migration").run(&connection).await?;
@@ -32,7 +66,56 @@ impl PostgresClient {
         Ok(Self(connection))
     }
 
+    pub fn pool(&self) -> &PgPool {
+        &self.0
+    }
+
     pub async fn acquire(&self) -> Result<PoolConnection<Postgres>, Error> {
         self.0.acquire().await
     }
 }
+
+#[cfg(all(test, not(target_os = "windows")))]
+mod test {
+    use hotshot_query_service::data_source::sql::testing::TmpDb;
+
+    use super::PostgresClient;
+    use crate::DatabaseOptions;
+
+    #[async_std::test]
+    async fn test_database_connection() {
+        let db = TmpDb::init().await;
+        let host = db.host();
+        let port = db.port();
+
+        let opts = DatabaseOptions {
+            url: None,
+            host: Some(host),
+            port: Some(port),
+            db_name: None,
+            username: Some("postgres".to_string()),
+            password: Some("password".to_string()),
+            max_connections: Some(100),
+            acquire_timeout: None,
+            migrations: true,
+        };
+
+        let client = PostgresClient::connect(opts)
+            .await
+            .expect("failed to connect to database");
+
+        let pool = client.pool();
+
+        sqlx::query("INSERT INTO test (str) VALUES ('testing');")
+            .execute(pool)
+            .await
+            .unwrap();
+
+        let result: i64 = sqlx::query_scalar("Select id from test where str = 'testing';")
+            .fetch_one(pool)
+            .await
+            .unwrap();
+
+        assert_eq!(result, 1);
+    }
+}

src/database.rs

@@ -16,10 +16,26 @@ pub struct Options {
 /// Arguments for establishing a database connection
 #[derive(Clone, Debug, Parser)]
 pub struct DatabaseOptions {
-    #[clap(long, env = "MARKETPLACE_SOLVER_DATABASE_URL")]
-    pub url: String,
+    // Postgres URL connection string
+    #[clap(long, env = "MARKETPLACE_SOLVER_POSTGRES_URL")]
+    pub url: Option<String>,
 
-    #[clap(long, env = "MARKETPLACE_SOLVER_DATABASE_MAX_CONNECTIONS")]
+    #[clap(long, env = "MARKETPLACE_SOLVER_POSTGRES_HOST")]
+    pub host: Option<String>,
+
+    #[clap(long, env = "MARKETPLACE_SOLVER_POSTGRES_PORT")]
+    pub port: Option<u16>,
+
+    #[clap(long, env = "MARKETPLACE_SOLVER_POSTGRES_DATABASE_NAME")]
+    pub db_name: Option<String>,
+
+    #[clap(long, env = "MARKETPLACE_SOLVER_POSTGRES_USERNAME")]
+    pub username: Option<String>,
+
+    #[clap(long, env = "MARKETPLACE_SOLVER_POSTGRES_PASSWORD")]
+    pub password: Option<String>,
+
+    #[clap(long, env = "MARKETPLACE_SOLVER_POSTGRES_MAX_CONNECTIONS")]
     pub max_connections: Option<u32>,
 
     #[clap(long,value_parser = parse_duration, env = "MARKETPLACE_SOLVER_DATABASE_ACQUIRE_TIMEOUT")]