Add PasswordProvider trait for dynamic Postgres credentials (#548)

* add password provider to postgres connector

* Fix clippy lint for PasswordProvider PR

* Use SecretString::clone in StaticPasswordProvider

---------

Co-authored-by: Phillip LeBlanc <phillip@leblanc.tech>

Author: alexanderbianchi

core/src/postgres.rs

@@ -1,6 +1,7 @@
 use crate::sql::arrow_sql_gen::statement::{
     CreateTableBuilder, Error as SqlGenError, IndexBuilder, InsertBuilder,
 };
+use crate::sql::db_connection_pool::dbconnection::postgresconn::PostgresPooledConnection;
 use crate::sql::db_connection_pool::{
     self,
     dbconnection::{postgresconn::PostgresConnection, DbConnection},
@@ -15,10 +16,7 @@ use arrow::{
     datatypes::{Schema, SchemaRef},
 };
 use async_trait::async_trait;
-use bb8_postgres::{
-    tokio_postgres::{types::ToSql, Transaction},
-    PostgresConnectionManager,
-};
+use bb8_postgres::tokio_postgres::{types::ToSql, Transaction};
 use datafusion::catalog::Session;
 use datafusion::sql::unparser::dialect::PostgreSqlDialect;
 use datafusion::{
@@ -29,7 +27,6 @@ use datafusion::{
     logical_expr::CreateExternalTable,
     sql::TableReference,
 };
-use postgres_native_tls::MakeTlsConnector;
 use snafu::prelude::*;
 use std::{collections::HashMap, sync::Arc};
 
@@ -47,15 +44,10 @@ use self::write::PostgresTableWriter;
 
 pub mod write;
 
-pub type DynPostgresConnectionPool = dyn DbConnectionPool<
-        bb8::PooledConnection<'static, PostgresConnectionManager<MakeTlsConnector>>,
-        &'static (dyn ToSql + Sync),
-    > + Send
-    + Sync;
-pub type DynPostgresConnection = dyn DbConnection<
-    bb8::PooledConnection<'static, PostgresConnectionManager<MakeTlsConnector>>,
-    &'static (dyn ToSql + Sync),
->;
+pub type DynPostgresConnectionPool =
+    dyn DbConnectionPool<PostgresPooledConnection, &'static (dyn ToSql + Sync)> + Send + Sync;
+pub type DynPostgresConnection =
+    dyn DbConnection<PostgresPooledConnection, &'static (dyn ToSql + Sync)>;
 
 #[derive(Debug, Snafu)]
 pub enum Error {

core/src/sql/db_connection_pool/dbconnection/postgresconn.rs

@@ -5,6 +5,7 @@ use std::sync::Arc;
 use crate::sql::arrow_sql_gen::postgres::rows_to_arrow;
 use crate::sql::arrow_sql_gen::postgres::schema::pg_data_type_to_arrow_type;
 use crate::sql::arrow_sql_gen::postgres::schema::ParseContext;
+use crate::sql::db_connection_pool::postgrespool::ConnectionManager;
 use crate::util::handle_unsupported_type_error;
 use crate::util::schema::SchemaValidator;
 use arrow::datatypes::Field;
@@ -13,14 +14,21 @@ use arrow::datatypes::SchemaRef;
 use arrow_schema::DataType;
 use async_stream::stream;
 use bb8_postgres::tokio_postgres::types::ToSql;
-use bb8_postgres::PostgresConnectionManager;
+
+/// A pooled Postgres connection obtained from a [`PostgresConnectionPool`](crate::sql::db_connection_pool::postgrespool::PostgresConnectionPool).
+///
+/// Dereferences to [`tokio_postgres::Client`](bb8_postgres::tokio_postgres::Client) for executing queries.
+// Defined here rather than in `postgrespool` to avoid a type-resolution cycle
+// between the two modules (postgrespool imports PostgresConnection, which uses
+// this alias).
+pub type PostgresPooledConnection = bb8::PooledConnection<'static, ConnectionManager>;
 use datafusion::error::DataFusionError;
 use datafusion::execution::SendableRecordBatchStream;
 use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
 use datafusion::sql::TableReference;
 use futures::stream;
 use futures::StreamExt;
-use postgres_native_tls::MakeTlsConnector;
+
 use snafu::prelude::*;
 
 use crate::UnsupportedTypeAction;
@@ -132,7 +140,7 @@ pub enum PostgresError {
 }
 
 pub struct PostgresConnection {
-    pub conn: bb8::PooledConnection<'static, PostgresConnectionManager<MakeTlsConnector>>,
+    pub conn: PostgresPooledConnection,
     unsupported_type_action: UnsupportedTypeAction,
 }
 
@@ -151,12 +159,7 @@ impl SchemaValidator for PostgresConnection {
     }
 }
 
-impl<'a>
-    DbConnection<
-        bb8::PooledConnection<'static, PostgresConnectionManager<MakeTlsConnector>>,
-        &'a (dyn ToSql + Sync),
-    > for PostgresConnection
-{
+impl<'a> DbConnection<PostgresPooledConnection, &'a (dyn ToSql + Sync)> for PostgresConnection {
     fn as_any(&self) -> &dyn Any {
         self
     }
@@ -167,26 +170,16 @@ impl<'a>
 
     fn as_async(
         &self,
-    ) -> Option<
-        &dyn AsyncDbConnection<
-            bb8::PooledConnection<'static, PostgresConnectionManager<MakeTlsConnector>>,
-            &'a (dyn ToSql + Sync),
-        >,
-    > {
+    ) -> Option<&dyn AsyncDbConnection<PostgresPooledConnection, &'a (dyn ToSql + Sync)>> {
         Some(self)
     }
 }
 
 #[async_trait::async_trait]
-impl<'a>
-    AsyncDbConnection<
-        bb8::PooledConnection<'static, PostgresConnectionManager<MakeTlsConnector>>,
-        &'a (dyn ToSql + Sync),
-    > for PostgresConnection
+impl<'a> AsyncDbConnection<PostgresPooledConnection, &'a (dyn ToSql + Sync)>
+    for PostgresConnection
 {
-    fn new(
-        conn: bb8::PooledConnection<'static, PostgresConnectionManager<MakeTlsConnector>>,
-    ) -> Self {
+    fn new(conn: PostgresPooledConnection) -> Self {
         PostgresConnection {
             conn,
             unsupported_type_action: UnsupportedTypeAction::default(),

core/src/sql/db_connection_pool/mod.rs

@@ -1,5 +1,6 @@
 use async_trait::async_trait;
 use dbconnection::DbConnection;
+use secrecy::SecretString;
 use std::sync::Arc;
 
 pub mod dbconnection;
@@ -23,6 +24,41 @@ pub mod sqlitepool;
 pub type Error = Box<dyn std::error::Error + Send + Sync>;
 type Result<T, E = Error> = std::result::Result<T, E>;
 
+/// A trait for providing passwords dynamically to database connection pools.
+///
+/// Implementations can fetch credentials from secret managers, IAM services,
+/// or other dynamic sources. Called each time a new connection is created in the pool.
+///
+/// Implementations that cache or rate-limit credentials should use interior
+/// mutability (e.g., `tokio::sync::RwLock`) since the trait requires `&self`.
+#[async_trait]
+pub trait PasswordProvider: Send + Sync {
+    /// Returns the current password/token for authentication.
+    /// Called each time a new connection is created in the pool.
+    async fn get_password(&self) -> Result<SecretString>;
+}
+
+/// A password provider that always returns the same static password.
+///
+/// This is the default provider used by [`PostgresConnectionPool::new()`](crate::sql::db_connection_pool::postgrespool::PostgresConnectionPool::new)
+/// when a `pass` parameter is supplied. It can also be used explicitly with
+/// [`new_with_password_provider()`](crate::sql::db_connection_pool::postgrespool::PostgresConnectionPool::new_with_password_provider).
+pub struct StaticPasswordProvider(SecretString);
+
+impl StaticPasswordProvider {
+    /// Creates a new `StaticPasswordProvider` with the given password.
+    pub fn new(password: SecretString) -> Self {
+        Self(password)
+    }
+}
+
+#[async_trait]
+impl PasswordProvider for StaticPasswordProvider {
+    async fn get_password(&self) -> Result<SecretString> {
+        Ok(self.0.clone())
+    }
+}
+
 /// Controls whether join pushdown is allowed, and under what conditions
 #[derive(Clone, Debug, PartialEq, Eq)]
 pub enum JoinPushDown {