connect to postgres via SSL (#993)

* connect to postgres via SSL

* Update app-server/src/db/mod.rs

Co-authored-by: ellipsis-dev[bot] <65095814+ellipsis-dev[bot]@users.noreply.github.com>

* fix build (don't require env vars to be set)

* fix regex

* replace .expect with anyhow error

* fix regex

---------

Co-authored-by: ellipsis-dev[bot] <65095814+ellipsis-dev[bot]@users.noreply.github.com>

Author: dinmukhamedm

app-server/Cargo.lock

@@ -45,7 +45,7 @@ serde = "1.0"
 serde_json = {version = "1.0.140", features = ["preserve_order"]}
 sha3 = "0.10.8"
 sodiumoxide = "0.2.7"
-sqlx = {version = "0.8", features = ["runtime-tokio", "postgres", "uuid", "json", "chrono", "bigdecimal"]}
+sqlx = {version = "0.8", features = ["runtime-tokio", "postgres", "uuid", "json", "chrono", "bigdecimal", "tls-rustls"]}
 thiserror = "2"
 tikv-jemallocator = "0.6"
 tokio = {version = "1.46", features = ["macros", "rt-multi-thread"]}

app-server/Cargo.toml

@@ -1,3 +1,6 @@
+use std::env;
+
+use regex::Regex;
 use sqlx::PgPool;
 
 pub mod datasets;
@@ -24,7 +27,80 @@ pub struct DB {
 }
 
 impl DB {
-    pub fn new(pool: PgPool) -> Self {
-        Self { pool }
+    pub async fn connect_from_env() -> anyhow::Result<Self> {
+        let options = get_pg_connect_options()?;
+        let pool = sqlx::postgres::PgPoolOptions::new()
+            .max_connections(
+                env::var("DATABASE_MAX_CONNECTIONS")
+                    .unwrap_or(String::from("10"))
+                    .parse::<u32>()
+                    .unwrap_or(10),
+            )
+            .connect_with(options)
+            .await?;
+        Ok(Self { pool })
+    }
+}
+
+fn get_pg_connect_options() -> anyhow::Result<sqlx::postgres::PgConnectOptions> {
+    let options = if let Ok(database_url) = env::var("DATABASE_URL") {
+        options_from_database_url(&database_url)?
+    } else {
+        options_from_database_env_vars()?
+    };
+
+    if let Ok(ssl_root_cert) = env::var("DATABASE_SSL_ROOT_CERT") {
+        Ok(options
+            .ssl_mode(sqlx::postgres::PgSslMode::VerifyFull)
+            .ssl_root_cert_from_pem(ssl_root_cert.into_bytes()))
+    } else {
+        Ok(options)
     }
 }
+
+fn options_from_database_url(
+    database_url: &str,
+) -> anyhow::Result<sqlx::postgres::PgConnectOptions> {
+    let re = Regex::new(r"^postgres(?:ql)?://([^:]+):([^@]+)@([^:]+):(\d*)/(.+)$").unwrap();
+    let caps = re
+        .captures(database_url)
+        .ok_or(anyhow::anyhow!("Invalid database URL"))?;
+    let username = caps.get(1).map(|m| m.as_str()).unwrap_or("postgres");
+    let password = caps
+        .get(2)
+        .ok_or(anyhow::anyhow!("Invalid database URL. Can't find password"))?
+        .as_str();
+    let host = caps
+        .get(3)
+        .ok_or(anyhow::anyhow!("Invalid database URL. Can't find host"))?
+        .as_str();
+    let port = caps
+        .get(4)
+        .and_then(|m| m.as_str().parse::<u16>().ok())
+        .unwrap_or(5432);
+    let database = caps.get(5).map(|m| m.as_str()).unwrap_or(username);
+    Ok(sqlx::postgres::PgConnectOptions::new()
+        .username(username)
+        .password(password)
+        .host(host)
+        .port(port)
+        .database(database))
+}
+
+fn options_from_database_env_vars() -> anyhow::Result<sqlx::postgres::PgConnectOptions> {
+    let username = env::var("DATABASE_USERNAME").unwrap_or(String::from("postgres"));
+    let password = env::var("DATABASE_PASSWORD")?;
+    let host = env::var("DATABASE_HOST")?;
+    let port = env::var("DATABASE_PORT")
+        .unwrap_or(String::from("5432"))
+        .parse::<u16>()
+        .unwrap_or(5432);
+    let database = env::var("DATABASE_DATABASE").unwrap_or(username.clone());
+
+    Ok(sqlx::postgres::PgConnectOptions::new()
+        .username(&username)
+        .password(&password)
+        .host(&host)
+        .port(port)
+        .database(&database))
+}

app-server/src/db/mod.rs

@@ -177,24 +177,8 @@ fn main() -> anyhow::Result<()> {
     let cache = Arc::new(cache);
 
     // === 2. Database ===
-    let db_url = env::var("DATABASE_URL").expect("DATABASE_URL must be set");
-
-    let max_connections = env::var("DATABASE_MAX_CONNECTIONS")
-        .unwrap_or(String::from("10"))
-        .parse()
-        .unwrap_or(10);
-
-    log::info!("Database max connections: {}", max_connections);
-
-    let pool = runtime_handle.block_on(async {
-        sqlx::postgres::PgPoolOptions::new()
-            .max_connections(max_connections)
-            .connect(&db_url)
-            .await
-            .unwrap()
-    });
-
-    let db = Arc::new(db::DB::new(pool));
+    let inner_db = runtime_handle.block_on(db::DB::connect_from_env())?;
+    let db = Arc::new(inner_db);
 
     // === 3. Message queues ===
     // Only enable RabbitMQ if it is a full build and RabbitMQ Feature (URL) is set
@@ -440,7 +424,7 @@ fn main() -> anyhow::Result<()> {
     let mq_for_http = queue.clone();
     let worker_tracker_for_http = worker_tracker.clone();
 
-    // == AWS config for S3 and Bedrock ==
+    // == AWS config for S3 ==
     let aws_sdk_config = runtime_handle.block_on(async {
         aws_config::defaults(BehaviorVersion::latest())
             .region(aws_config::Region::new(