Add IO runtime handle to postgres connector (#549)

Author: alexanderbianchi

core/src/sql/db_connection_pool/postgrespool.rs

@@ -11,6 +11,7 @@ use native_tls::{Certificate, TlsConnector};
 use postgres_native_tls::MakeTlsConnector;
 use secrecy::{ExposeSecret, SecretBox, SecretString};
 use snafu::{prelude::*, ResultExt};
+use tokio::runtime::Handle;
 use tokio_postgres;
 
 use super::{
@@ -79,6 +80,9 @@ pub enum Error {
     PasswordProviderError {
         source: Box<dyn std::error::Error + Send + Sync>,
     },
+
+    #[snafu(display("Task failed to execute on IO runtime.\n{source}"))]
+    IoRuntimeError { source: tokio::task::JoinError },
 }
 
 pub type Result<T, E = Error> = std::result::Result<T, E>;
@@ -187,6 +191,7 @@ pub struct PostgresConnectionPool {
     pool: Arc<bb8::Pool<ConnectionManager>>,
     join_push_down: JoinPushDown,
     unsupported_type_action: UnsupportedTypeAction,
+    io_handle: Option<Handle>,
 }
 
 impl PostgresConnectionPool {
@@ -380,6 +385,7 @@ impl PostgresConnectionPool {
             pool: Arc::new(pool),
             join_push_down,
             unsupported_type_action: UnsupportedTypeAction::default(),
+            io_handle: None,
         })
     }
 
@@ -390,14 +396,28 @@ impl PostgresConnectionPool {
         self
     }
 
+    /// Route all Postgres connection background tasks to a dedicated IO runtime.
+    #[must_use]
+    pub fn with_io_runtime(mut self, handle: Handle) -> Self {
+        self.io_handle = Some(handle);
+        self
+    }
+
     /// Returns a direct connection to the underlying database.
     ///
     /// # Errors
     ///
     /// Returns an error if there is a problem creating the connection pool.
     pub async fn connect_direct(&self) -> super::Result<PostgresConnection> {
         let pool = Arc::clone(&self.pool);
-        let conn = pool.get_owned().await.map_err(map_pool_run_error)?;
+        let conn = if let Some(handle) = &self.io_handle {
+            handle
+                .spawn(async move { pool.get_owned().await.map_err(map_pool_run_error) })
+                .await
+                .context(IoRuntimeSnafu)??
+        } else {
+            pool.get_owned().await.map_err(map_pool_run_error)?
+        };
         Ok(PostgresConnection::new(conn))
     }
 }
@@ -581,8 +601,15 @@ impl
         >,
     > {
         let pool = Arc::clone(&self.pool);
-        let get_conn = async || pool.get_owned().await.map_err(map_pool_run_error);
-        let conn = run_async_with_tokio(get_conn).await?;
+        let conn = if let Some(handle) = &self.io_handle {
+            handle
+                .spawn(async move { pool.get_owned().await.map_err(map_pool_run_error) })
+                .await
+                .context(IoRuntimeSnafu)??
+        } else {
+            let get_conn = async || pool.get_owned().await.map_err(map_pool_run_error);
+            run_async_with_tokio(get_conn).await?
+        };
         Ok(Box::new(
             PostgresConnection::new(conn)
                 .with_unsupported_type_action(self.unsupported_type_action),

core/tests/postgres/mod.rs

@@ -453,3 +453,37 @@ async fn arrow_postgres_one_way(
 
     assert_eq!(record_batch[0], expected_record);
 }
+
+#[rstest]
+#[test_log::test(tokio::test)]
+async fn test_postgres_io_runtime_segregation(container_manager: &Mutex<ContainerManager>) {
+    let mut container_manager = container_manager.lock().await;
+    if !container_manager.claimed {
+        container_manager.claimed = true;
+        start_container(&mut container_manager).await;
+    }
+
+    // Create a separate IO runtime
+    let io_runtime = tokio::runtime::Builder::new_multi_thread()
+        .worker_threads(2)
+        .enable_all()
+        .build()
+        .expect("IO runtime should be created");
+
+    let pool = common::get_postgres_connection_pool(container_manager.port)
+        .await
+        .expect("pool created")
+        .with_io_runtime(io_runtime.handle().clone());
+
+    // Verify the pool works through the IO runtime
+    let sqltable_pool: Arc<DynPostgresConnectionPool> = Arc::new(pool);
+    let conn = sqltable_pool.connect().await.expect("connect should work");
+    let async_conn = conn.as_async().expect("should be async connection");
+    // Execute a simple query to confirm IO runtime is functional
+    let stream = async_conn
+        .query_arrow("SELECT 1 AS val", &[], None)
+        .await
+        .expect("query should work");
+    let batches: Vec<_> = futures::StreamExt::collect(stream).await;
+    assert!(!batches.is_empty(), "should return results via IO runtime");
+}