From 4a5bcd232b9ef70d1262ce73fba917a74a546ff3 Mon Sep 17 00:00:00 2001
From: Marko Malenic <mmalenic1@gmail.com>
Date: Mon, 1 Jul 2024 19:23:08 +1000
Subject: [PATCH] fix(filemanager): avoid using to_url_lossy when loading
 database credentials

---
 .../filemanager/src/database/mod.rs           | 27 +++++++------------
 .../filemanager/src/handlers/aws.rs           |  6 +++--
 2 files changed, 14 insertions(+), 19 deletions(-)

diff --git a/lib/workload/stateless/stacks/filemanager/filemanager/src/database/mod.rs b/lib/workload/stateless/stacks/filemanager/filemanager/src/database/mod.rs
index cde002207..75edad470 100644
--- a/lib/workload/stateless/stacks/filemanager/filemanager/src/database/mod.rs
+++ b/lib/workload/stateless/stacks/filemanager/filemanager/src/database/mod.rs
@@ -5,9 +5,9 @@ use crate::database::aws::ingester::Ingester;
 use crate::database::aws::ingester_paired::IngesterPaired;
 use crate::env::Config;
 use async_trait::async_trait;
-use sea_orm::{ConnectOptions, Database, DatabaseConnection, SqlxPostgresConnector};
+use sea_orm::{DatabaseConnection, SqlxPostgresConnector};
 use sqlx::postgres::PgConnectOptions;
-use sqlx::{ConnectOptions as SqlxConnectOptions, PgPool};
+use sqlx::PgPool;
 use tracing::debug;
 
 use crate::error::Result;
@@ -49,7 +49,7 @@ impl Client {
         generator: Option<impl CredentialGenerator>,
         config: &Config,
     ) -> Result<Self> {
-        Ok(Self::new(Self::create_pool(generator, config).await?))
+        Ok(Self::from_pool(Self::create_pool(generator, config).await?))
     }
 
     /// Create a database connection pool using credential loading logic defined in
@@ -57,20 +57,8 @@ impl Client {
     pub async fn create_pool(
         generator: Option<impl CredentialGenerator>,
         config: &Config,
-    ) -> Result<DatabaseConnection> {
-        Ok(Database::connect(Self::connect_options(generator, config).await?).await?)
-    }
-
-    /// Create database connect options using a series of credential loading logic.
-    pub async fn connect_options(
-        generator: Option<impl CredentialGenerator>,
-        config: &Config,
-    ) -> Result<ConnectOptions> {
-        Ok(ConnectOptions::new(
-            Self::pg_connect_options(generator, config)
-                .await?
-                .to_url_lossy(),
-        ))
+    ) -> Result<PgPool> {
+        Ok(PgPool::connect_with(Self::pg_connect_options(generator, config).await?).await?)
     }
 
     /// Create database connect options using a series of credential loading logic.
@@ -117,6 +105,11 @@ impl Client {
     pub fn connection_ref(&self) -> &DatabaseConnection {
         &self.connection
     }
+
+    /// Get the inner database connection.
+    pub fn into_inner(self) -> DatabaseConnection {
+        self.connection
+    }
 }
 
 #[async_trait]
diff --git a/lib/workload/stateless/stacks/filemanager/filemanager/src/handlers/aws.rs b/lib/workload/stateless/stacks/filemanager/filemanager/src/handlers/aws.rs
index f2e9cab9d..4ebaa01b0 100644
--- a/lib/workload/stateless/stacks/filemanager/filemanager/src/handlers/aws.rs
+++ b/lib/workload/stateless/stacks/filemanager/filemanager/src/handlers/aws.rs
@@ -169,11 +169,13 @@ pub async fn ingest_s3_inventory(
 
 /// Create a postgres database pool using an IAM credential generator.
 pub async fn create_database_pool(env_config: &EnvConfig) -> Result<DatabaseConnection, Error> {
-    Ok(Client::create_pool(
+    let client = Client::from_generator(
         Some(IamGeneratorBuilder::default().build(env_config).await?),
         env_config,
     )
-    .await?)
+    .await?;
+
+    Ok(client.into_inner())
 }
 
 /// Update connection options with new credentials.