Refactor statement timeout to be truly per-session using ClientInfo metadata

- Remove statement_timeout field from DfSessionService (was application-scoped)
- Store timeout in ClientInfo::metadata() using METADATA_STATEMENT_TIMEOUT key
- Add helper functions get_statement_timeout() and set_statement_timeout()
- Update SET/SHOW statement_timeout handlers to use client metadata
- Update query execution logic to read timeout from client session
- Add comprehensive MockClient for testing session-specific behavior
- Now each PostgreSQL session has its own independent timeout setting

This follows the PostgreSQL standard where statement_timeout is a session variable,
not a server-wide configuration.

Author: iPeluwa

datafusion-postgres/src/handlers.rs

@@ -31,6 +31,9 @@ use tokio::sync::Mutex;
 use arrow_pg::datatypes::df;
 use arrow_pg::datatypes::{arrow_schema_to_pg_fields, into_pg_type};
 
+// Metadata keys for session-level settings
+const METADATA_STATEMENT_TIMEOUT: &str = "statement_timeout_ms";
+
 /// Simple startup handler that does no authentication
 /// For production, use DfAuthSource with proper pgwire authentication handlers
 pub struct SimpleStartupHandler;
@@ -71,7 +74,6 @@ pub struct DfSessionService {
     timezone: Arc<Mutex<String>>,
     auth_manager: Arc<AuthManager>,
     sql_rewrite_rules: Vec<Arc<dyn SqlStatementRewriteRule>>,
-    statement_timeout: Arc<Mutex<Option<std::time::Duration>>>,
 }
 
 impl DfSessionService {
@@ -98,7 +100,31 @@ impl DfSessionService {
             timezone: Arc::new(Mutex::new("UTC".to_string())),
             auth_manager,
             sql_rewrite_rules,
-            statement_timeout: Arc::new(Mutex::new(None)),
+        }
+    }
+
+    /// Get statement timeout from client metadata
+    fn get_statement_timeout<C>(client: &C) -> Option<std::time::Duration>
+    where
+        C: ClientInfo,
+    {
+        client
+            .metadata()
+            .get(METADATA_STATEMENT_TIMEOUT)
+            .and_then(|s| s.parse::<u64>().ok())
+            .map(std::time::Duration::from_millis)
+    }
+
+    /// Set statement timeout in client metadata
+    fn set_statement_timeout<C>(client: &mut C, timeout: Option<std::time::Duration>)
+    where
+        C: ClientInfo,
+    {
+        let metadata = client.metadata_mut();
+        if let Some(duration) = timeout {
+            metadata.insert(METADATA_STATEMENT_TIMEOUT.to_string(), duration.as_millis().to_string());
+        } else {
+            metadata.remove(METADATA_STATEMENT_TIMEOUT);
         }
     }
 
@@ -196,10 +222,14 @@ impl DfSessionService {
         Ok(QueryResponse::new(Arc::new(fields), Box::pin(row_stream)))
     }
 
-    async fn try_respond_set_statements<'a>(
+    async fn try_respond_set_statements<'a, C>(
         &self,
+        client: &mut C,
         query_lower: &str,
-    ) -> PgWireResult<Option<Response<'a>>> {
+    ) -> PgWireResult<Option<Response<'a>>>
+    where
+        C: ClientInfo,
+    {
         if query_lower.starts_with("set") {
             if query_lower.starts_with("set time zone") {
                 let parts: Vec<&str> = query_lower.split_whitespace().collect();
@@ -221,10 +251,9 @@ impl DfSessionService {
                 let parts: Vec<&str> = query_lower.split_whitespace().collect();
                 if parts.len() >= 3 {
                     let timeout_str = parts[2].trim_matches('"').trim_matches('\'');
-                    let mut statement_timeout = self.statement_timeout.lock().await;
 
-                    if timeout_str == "0" || timeout_str.is_empty() {
-                        *statement_timeout = None;
+                    let timeout = if timeout_str == "0" || timeout_str.is_empty() {
+                        None
                     } else {
                         // Parse timeout value (supports ms, s, min formats)
                         let timeout_ms = if timeout_str.ends_with("ms") {
@@ -245,14 +274,12 @@ impl DfSessionService {
                         };
 
                         match timeout_ms {
-                            Ok(ms) if ms > 0 => {
-                                *statement_timeout = Some(std::time::Duration::from_millis(ms));
-                            }
-                            _ => {
-                                *statement_timeout = None;
-                            }
+                            Ok(ms) if ms > 0 => Some(std::time::Duration::from_millis(ms)),
+                            _ => None,
                         }
-                    }
+                    };
+
+                    Self::set_statement_timeout(client, timeout);
                     Ok(Some(Response::Execution(Tag::new("SET"))))
                 } else {
                     Err(PgWireError::UserError(Box::new(
@@ -322,10 +349,14 @@ impl DfSessionService {
         }
     }
 
-    async fn try_respond_show_statements<'a>(
+    async fn try_respond_show_statements<'a, C>(
         &self,
+        client: &C,
         query_lower: &str,
-    ) -> PgWireResult<Option<Response<'a>>> {
+    ) -> PgWireResult<Option<Response<'a>>>
+    where
+        C: ClientInfo,
+    {
         if query_lower.starts_with("show ") {
             match query_lower.strip_suffix(";").unwrap_or(query_lower) {
                 "show time zone" => {
@@ -354,7 +385,7 @@ impl DfSessionService {
                     Ok(Some(Response::Query(resp)))
                 }
                 "show statement_timeout" => {
-                    let timeout = *self.statement_timeout.lock().await;
+                    let timeout = Self::get_statement_timeout(client);
                     let timeout_str = match timeout {
                         Some(duration) => format!("{}ms", duration.as_millis()),
                         None => "0".to_string(),
@@ -408,7 +439,7 @@ impl SimpleQueryHandler for DfSessionService {
             self.check_query_permission(client, &query).await?;
         }
 
-        if let Some(resp) = self.try_respond_set_statements(&query_lower).await? {
+        if let Some(resp) = self.try_respond_set_statements(client, &query_lower).await? {
             return Ok(vec![resp]);
         }
 
@@ -419,7 +450,7 @@ impl SimpleQueryHandler for DfSessionService {
             return Ok(vec![resp]);
         }
 
-        if let Some(resp) = self.try_respond_show_statements(&query_lower).await? {
+        if let Some(resp) = self.try_respond_show_statements(client, &query_lower).await? {
             return Ok(vec![resp]);
         }
 
@@ -436,7 +467,7 @@ impl SimpleQueryHandler for DfSessionService {
         }
 
         let df_result = {
-            let timeout = *self.statement_timeout.lock().await;
+            let timeout = Self::get_statement_timeout(client);
             if let Some(timeout_duration) = timeout {
                 tokio::time::timeout(timeout_duration, self.session_context.sql(&query))
                     .await
@@ -568,7 +599,7 @@ impl ExtendedQueryHandler for DfSessionService {
                 .await?;
         }
 
-        if let Some(resp) = self.try_respond_set_statements(&query).await? {
+        if let Some(resp) = self.try_respond_set_statements(client, &query).await? {
             return Ok(resp);
         }
 
@@ -579,7 +610,7 @@ impl ExtendedQueryHandler for DfSessionService {
             return Ok(resp);
         }
 
-        if let Some(resp) = self.try_respond_show_statements(&query).await? {
+        if let Some(resp) = self.try_respond_show_statements(client, &query).await? {
             return Ok(resp);
         }
 
@@ -613,7 +644,7 @@ impl ExtendedQueryHandler for DfSessionService {
             .map_err(|e| PgWireError::ApiError(Box::new(e)))?;
 
         let dataframe = {
-            let timeout = *self.statement_timeout.lock().await;
+            let timeout = Self::get_statement_timeout(client);
             if let Some(timeout_duration) = timeout {
                 tokio::time::timeout(
                     timeout_duration,
@@ -690,28 +721,88 @@ mod tests {
     use super::*;
     use crate::auth::AuthManager;
     use datafusion::prelude::SessionContext;
+    use std::collections::HashMap;
     use std::time::Duration;
 
+    struct MockClient {
+        metadata: HashMap<String, String>,
+    }
+
+    impl MockClient {
+        fn new() -> Self {
+            Self {
+                metadata: HashMap::new(),
+            }
+        }
+    }
+
+    impl ClientInfo for MockClient {
+        fn socket_addr(&self) -> std::net::SocketAddr {
+            "127.0.0.1:5432".parse().unwrap()
+        }
+
+        fn is_secure(&self) -> bool {
+            false
+        }
+
+        fn protocol_version(&self) -> pgwire::messages::ProtocolVersion {
+            pgwire::messages::ProtocolVersion::PROTOCOL3_0
+        }
+
+        fn set_protocol_version(&mut self, _version: pgwire::messages::ProtocolVersion) {}
+
+        fn pid_and_secret_key(&self) -> (i32, pgwire::messages::startup::SecretKey) {
+            (0, pgwire::messages::startup::SecretKey::I32(0))
+        }
+
+        fn set_pid_and_secret_key(&mut self, _pid: i32, _secret_key: pgwire::messages::startup::SecretKey) {}
+
+        fn state(&self) -> pgwire::api::PgWireConnectionState {
+            pgwire::api::PgWireConnectionState::ReadyForQuery
+        }
+
+        fn set_state(&mut self, _new_state: pgwire::api::PgWireConnectionState) {}
+
+        fn transaction_status(&self) -> pgwire::messages::response::TransactionStatus {
+            pgwire::messages::response::TransactionStatus::Idle
+        }
+
+        fn set_transaction_status(&mut self, _new_status: pgwire::messages::response::TransactionStatus) {}
+
+        fn metadata(&self) -> &HashMap<String, String> {
+            &self.metadata
+        }
+
+        fn metadata_mut(&mut self) -> &mut HashMap<String, String> {
+            &mut self.metadata
+        }
+
+        fn client_certificates<'a>(&self) -> Option<&[rustls_pki_types::CertificateDer<'a>]> {
+            None
+        }
+    }
+
     #[tokio::test]
     async fn test_statement_timeout_set_and_show() {
         let session_context = Arc::new(SessionContext::new());
         let auth_manager = Arc::new(AuthManager::new());
         let service = DfSessionService::new(session_context, auth_manager);
+        let mut client = MockClient::new();
 
         // Test setting timeout to 5000ms
         let set_response = service
-            .try_respond_set_statements("set statement_timeout '5000ms'")
+            .try_respond_set_statements(&mut client, "set statement_timeout '5000ms'")
             .await
             .unwrap();
         assert!(set_response.is_some());
 
-        // Verify the timeout was set
-        let timeout = *service.statement_timeout.lock().await;
+        // Verify the timeout was set in client metadata
+        let timeout = DfSessionService::get_statement_timeout(&client);
         assert_eq!(timeout, Some(Duration::from_millis(5000)));
 
         // Test SHOW statement_timeout
         let show_response = service
-            .try_respond_show_statements("show statement_timeout")
+            .try_respond_show_statements(&client, "show statement_timeout")
             .await
             .unwrap();
         assert!(show_response.is_some());
@@ -722,20 +813,21 @@ mod tests {
         let session_context = Arc::new(SessionContext::new());
         let auth_manager = Arc::new(AuthManager::new());
         let service = DfSessionService::new(session_context, auth_manager);
+        let mut client = MockClient::new();
 
         // Set timeout first
         service
-            .try_respond_set_statements("set statement_timeout '1000ms'")
+            .try_respond_set_statements(&mut client, "set statement_timeout '1000ms'")
             .await
             .unwrap();
 
         // Disable timeout with 0
         service
-            .try_respond_set_statements("set statement_timeout '0'")
+            .try_respond_set_statements(&mut client, "set statement_timeout '0'")
             .await
             .unwrap();
 
-        let timeout = *service.statement_timeout.lock().await;
+        let timeout = DfSessionService::get_statement_timeout(&client);
         assert_eq!(timeout, None);
     }
 }