refactor(drainer, router): KMS decrypt database password when kms feature is enabled (#733)

crates/router/src/configs/defaults.rs

@@ -15,12 +15,15 @@ impl Default for super::settings::Database {
     fn default() -> Self {
         Self {
             username: String::new(),
+            #[cfg(not(feature = "kms"))]
             password: String::new(),
             host: "localhost".into(),
             port: 5432,
             dbname: String::new(),
             pool_size: 5,
             connection_timeout: 10,
+            #[cfg(feature = "kms")]
+            kms_encrypted_password: String::new(),
         }
     }
 }

crates/router/src/configs/settings.rs

@@ -223,12 +223,15 @@ pub struct Server {
 #[serde(default)]
 pub struct Database {
     pub username: String,
+    #[cfg(not(feature = "kms"))]
     pub password: String,
     pub host: String,
     pub port: u16,
     pub dbname: String,
     pub pool_size: u32,
     pub connection_timeout: u64,
+    #[cfg(feature = "kms")]
+    pub kms_encrypted_password: String,
 }
 
 #[derive(Debug, Deserialize, Clone)]

crates/router/src/configs/validations.rs

@@ -61,23 +61,35 @@ impl super::settings::Database {
             ))
         })?;
 
+        when(self.dbname.is_default_or_empty(), || {
+            Err(ApplicationError::InvalidConfigurationValueError(
+                "database name must not be empty".into(),
+            ))
+        })?;
+
         when(self.username.is_default_or_empty(), || {
             Err(ApplicationError::InvalidConfigurationValueError(
                 "database user username must not be empty".into(),
             ))
         })?;
 
-        when(self.password.is_default_or_empty(), || {
-            Err(ApplicationError::InvalidConfigurationValueError(
-                "database user password must not be empty".into(),
-            ))
-        })?;
+        #[cfg(not(feature = "kms"))]
+        {
+            when(self.password.is_default_or_empty(), || {
+                Err(ApplicationError::InvalidConfigurationValueError(
+                    "database user password must not be empty".into(),
+                ))
+            })
+        }
 
-        when(self.dbname.is_default_or_empty(), || {
-            Err(ApplicationError::InvalidConfigurationValueError(
-                "database name must not be empty".into(),
-            ))
-        })
+        #[cfg(feature = "kms")]
+        {
+            when(self.kms_encrypted_password.is_default_or_empty(), || {
+                Err(ApplicationError::InvalidConfigurationValueError(
+                    "database KMS encrypted password must not be empty".into(),
+                ))
+            })
+        }
     }
 }
 

crates/router/src/connection.rs

@@ -2,6 +2,8 @@ use async_bb8_diesel::{AsyncConnection, ConnectionError};
 use bb8::{CustomizeConnection, PooledConnection};
 use diesel::PgConnection;
 use error_stack::{IntoReport, ResultExt};
+#[cfg(feature = "kms")]
+use external_services::kms;
 
 use crate::{configs::settings::Database, errors};
 
@@ -37,10 +39,24 @@ pub async fn redis_connection(
 }
 
 #[allow(clippy::expect_used)]
-pub async fn diesel_make_pg_pool(database: &Database, test_transaction: bool) -> PgPool {
+pub async fn diesel_make_pg_pool(
+    database: &Database,
+    test_transaction: bool,
+    #[cfg(feature = "kms")] kms_config: &kms::KmsConfig,
+) -> PgPool {
+    #[cfg(feature = "kms")]
+    let password = kms::get_kms_client(kms_config)
+        .await
+        .decrypt(&database.kms_encrypted_password)
+        .await
+        .expect("Failed to KMS decrypt database password");
+
+    #[cfg(not(feature = "kms"))]
+    let password = &database.password;
+
     let database_url = format!(
         "postgres://{}:{}@{}:{}/{}",
-        database.username, database.password, database.host, database.port, database.dbname
+        database.username, password, database.host, database.port, database.dbname
     );
     let manager = async_bb8_diesel::ConnectionManager::<PgConnection>::new(database_url);
     let mut pool = bb8::Pool::builder()

crates/router/src/services.rs

@@ -109,9 +109,21 @@ impl Store {
         });
 
         Self {
-            master_pool: diesel_make_pg_pool(&config.master_database, test_transaction).await,
+            master_pool: diesel_make_pg_pool(
+                &config.master_database,
+                test_transaction,
+                #[cfg(feature = "kms")]
+                &config.kms,
+            )
+            .await,
             #[cfg(feature = "olap")]
-            replica_pool: diesel_make_pg_pool(&config.replica_database, test_transaction).await,
+            replica_pool: diesel_make_pg_pool(
+                &config.replica_database,
+                test_transaction,
+                #[cfg(feature = "kms")]
+                &config.kms,
+            )
+            .await,
             redis_conn,
             #[cfg(feature = "kv_store")]
             config: StoreConfig {