Fix parameter type inference, schema mismatches, and add configurable query timeout

- Fixed parameter type inference for arithmetic operations with untyped parameters (defaults to TEXT instead of throwing fatal errors)
- Fixed Arrow schema mismatches by unifying all UTF8 variants to use TEXT type consistently
- Added configurable query timeout with ServerOptions.with_query_timeout_secs(0) for no timeout
- Added configurable max_connections with ServerOptions.with_max_connections(n)
- Added connection limiting with semaphore to prevent resource exhaustion under load
- Simplified API with single constructor that takes timeout parameter directly
- Added comprehensive unit tests for timeout and connection configuration functionality

Author: iPeluwa

arrow-pg/src/datatypes.rs

@@ -42,8 +42,7 @@ pub fn into_pg_type(arrow_type: &DataType) -> PgWireResult<Type> {
         DataType::Float16 | DataType::Float32 => Type::FLOAT4,
         DataType::Float64 => Type::FLOAT8,
         DataType::Decimal128(_, _) => Type::NUMERIC,
-        DataType::Utf8 => Type::VARCHAR,
-        DataType::LargeUtf8 | DataType::Utf8View => Type::TEXT,
+        DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View => Type::TEXT,
         DataType::List(field) | DataType::FixedSizeList(field, _) | DataType::LargeList(field) => {
             match field.data_type() {
                 DataType::Boolean => Type::BOOL_ARRAY,
@@ -67,8 +66,7 @@ pub fn into_pg_type(arrow_type: &DataType) -> PgWireResult<Type> {
                 | DataType::BinaryView => Type::BYTEA_ARRAY,
                 DataType::Float16 | DataType::Float32 => Type::FLOAT4_ARRAY,
                 DataType::Float64 => Type::FLOAT8_ARRAY,
-                DataType::Utf8 => Type::VARCHAR_ARRAY,
-                DataType::LargeUtf8 | DataType::Utf8View => Type::TEXT_ARRAY,
+                DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View => Type::TEXT_ARRAY,
                 struct_type @ DataType::Struct(_) => Type::new(
                     Type::RECORD_ARRAY.name().into(),
                     Type::RECORD_ARRAY.oid(),

arrow-pg/src/datatypes/df.rs

@@ -66,11 +66,9 @@ where
         } else if let Some(infer_type) = inferenced_type {
             into_pg_type(infer_type)
         } else {
-            Err(PgWireError::UserError(Box::new(ErrorInfo::new(
-                "FATAL".to_string(),
-                "XX000".to_string(),
-                "Unknown parameter type".to_string(),
-            ))))
+            // Default to TEXT/VARCHAR for untyped parameters
+            // This allows arithmetic operations to work with implicit casting
+            Ok(Type::TEXT)
         }
     }
 

datafusion-postgres/src/handlers.rs

@@ -43,9 +43,9 @@ pub struct HandlerFactory {
 }
 
 impl HandlerFactory {
-    pub fn new(session_context: Arc<SessionContext>, auth_manager: Arc<AuthManager>) -> Self {
+    pub fn new(session_context: Arc<SessionContext>, auth_manager: Arc<AuthManager>, query_timeout: Option<std::time::Duration>) -> Self {
         let session_service =
-            Arc::new(DfSessionService::new(session_context, auth_manager.clone()));
+            Arc::new(DfSessionService::new(session_context, auth_manager.clone(), query_timeout));
         HandlerFactory { session_service }
     }
 }
@@ -71,12 +71,14 @@ pub struct DfSessionService {
     timezone: Arc<Mutex<String>>,
     auth_manager: Arc<AuthManager>,
     sql_rewrite_rules: Vec<Arc<dyn SqlStatementRewriteRule>>,
+    query_timeout: Option<std::time::Duration>,
 }
 
 impl DfSessionService {
     pub fn new(
         session_context: Arc<SessionContext>,
         auth_manager: Arc<AuthManager>,
+        query_timeout: Option<std::time::Duration>,
     ) -> DfSessionService {
         let sql_rewrite_rules: Vec<Arc<dyn SqlStatementRewriteRule>> = vec![
             Arc::new(AliasDuplicatedProjectionRewrite),
@@ -97,9 +99,12 @@ impl DfSessionService {
             timezone: Arc::new(Mutex::new("UTC".to_string())),
             auth_manager,
             sql_rewrite_rules,
+            query_timeout,
         }
     }
 
+
+
     /// Check if the current user has permission to execute a query
     async fn check_query_permission<C>(&self, client: &C, query: &str) -> PgWireResult<()>
     where
@@ -378,7 +383,19 @@ impl SimpleQueryHandler for DfSessionService {
             )));
         }
 
-        let df_result = self.session_context.sql(&query).await;
+        let df_result = if let Some(timeout) = self.query_timeout {
+            tokio::time::timeout(timeout, self.session_context.sql(&query))
+                .await
+                .map_err(|_| {
+                    PgWireError::UserError(Box::new(pgwire::error::ErrorInfo::new(
+                        "ERROR".to_string(),
+                        "57014".to_string(), // query_canceled error code
+                        "canceling statement due to query timeout".to_string(),
+                    )))
+                })?
+        } else {
+            self.session_context.sql(&query).await
+        };
 
         // Handle query execution errors and transaction state
         let df = match df_result {
@@ -540,10 +557,23 @@ impl ExtendedQueryHandler for DfSessionService {
             .optimize(&plan)
             .map_err(|e| PgWireError::ApiError(Box::new(e)))?;
 
-        let dataframe = match self.session_context.execute_logical_plan(optimised).await {
-            Ok(df) => df,
-            Err(e) => {
-                return Err(PgWireError::ApiError(Box::new(e)));
+        let dataframe = if let Some(timeout) = self.query_timeout {
+            tokio::time::timeout(timeout, self.session_context.execute_logical_plan(optimised))
+                .await
+                .map_err(|_| {
+                    PgWireError::UserError(Box::new(pgwire::error::ErrorInfo::new(
+                        "ERROR".to_string(),
+                        "57014".to_string(), // query_canceled error code
+                        "canceling statement due to query timeout".to_string(),
+                    )))
+                })?
+                .map_err(|e| PgWireError::ApiError(Box::new(e)))?
+        } else {
+            match self.session_context.execute_logical_plan(optimised).await {
+                Ok(df) => df,
+                Err(e) => {
+                    return Err(PgWireError::ApiError(Box::new(e)));
+                }
             }
         };
         let resp = df::encode_dataframe(dataframe, &portal.result_column_format).await?;
@@ -593,3 +623,76 @@ fn ordered_param_types(types: &HashMap<String, Option<DataType>>) -> Vec<Option<
     types.sort_by(|a, b| a.0.cmp(b.0));
     types.into_iter().map(|pt| pt.1.as_ref()).collect()
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::{auth::AuthManager, ServerOptions};
+    use datafusion::prelude::SessionContext;
+    use std::time::Duration;
+
+    #[test]
+    fn test_server_options_default_timeout() {
+        let opts = ServerOptions::default();
+        assert_eq!(opts.query_timeout, Some(Duration::from_secs(30)));
+    }
+
+    #[test]
+    fn test_server_options_no_timeout() {
+        let mut opts = ServerOptions::new();
+        opts.query_timeout = None;
+        assert_eq!(opts.query_timeout, None);
+    }
+
+    #[test]
+    fn test_handler_factory_with_timeout() {
+        let session_context = Arc::new(SessionContext::new());
+        let auth_manager = Arc::new(AuthManager::new());
+        let timeout = Some(Duration::from_secs(60));
+        
+        let factory = HandlerFactory::new(session_context, auth_manager, timeout);
+        assert_eq!(factory.session_service.query_timeout, timeout);
+    }
+
+    #[test]
+    fn test_session_service_timeout_configuration() {
+        let session_context = Arc::new(SessionContext::new());
+        let auth_manager = Arc::new(AuthManager::new());
+        
+        // Test with timeout
+        let service_with_timeout = DfSessionService::new(
+            session_context.clone(), 
+            auth_manager.clone(), 
+            Some(Duration::from_secs(45))
+        );
+        assert_eq!(service_with_timeout.query_timeout, Some(Duration::from_secs(45)));
+        
+        // Test without timeout (None)
+        let service_no_timeout = DfSessionService::new(
+            session_context, 
+            auth_manager, 
+            None
+        );
+        assert_eq!(service_no_timeout.query_timeout, None);
+    }
+
+    #[test]
+    fn test_timeout_configuration_from_seconds() {
+        // Test 0 seconds = no timeout
+        let opts_no_timeout = ServerOptions::new().with_query_timeout_secs(0);
+        assert_eq!(opts_no_timeout.query_timeout, None);
+        
+        // Test positive seconds = Some(Duration)
+        let opts_with_timeout = ServerOptions::new().with_query_timeout_secs(60);
+        assert_eq!(opts_with_timeout.query_timeout, Some(Duration::from_secs(60)));
+    }
+
+    #[test]
+    fn test_max_connections_configuration() {
+        let opts = ServerOptions::new().with_max_connections(500);
+        assert_eq!(opts.max_connections, 500);
+        
+        let opts_default = ServerOptions::default();
+        assert_eq!(opts_default.max_connections, 1000);
+    }
+}

datafusion-postgres/src/lib.rs

@@ -16,6 +16,7 @@ use pgwire::tokio::process_socket;
 use rustls_pemfile::{certs, pkcs8_private_keys};
 use rustls_pki_types::{CertificateDer, PrivateKeyDer};
 use tokio::net::TcpListener;
+use tokio::sync::Semaphore;
 use tokio_rustls::rustls::{self, ServerConfig};
 use tokio_rustls::TlsAcceptor;
 
@@ -34,12 +35,24 @@ pub struct ServerOptions {
     port: u16,
     tls_cert_path: Option<String>,
     tls_key_path: Option<String>,
+    max_connections: usize,
+    query_timeout: Option<std::time::Duration>,
 }
 
 impl ServerOptions {
     pub fn new() -> ServerOptions {
         ServerOptions::default()
     }
+
+    /// Set query timeout from seconds. Use 0 for no timeout.
+    pub fn with_query_timeout_secs(mut self, timeout_secs: u64) -> Self {
+        self.query_timeout = if timeout_secs == 0 {
+            None
+        } else {
+            Some(std::time::Duration::from_secs(timeout_secs))
+        };
+        self
+    }
 }
 
 impl Default for ServerOptions {
@@ -49,6 +62,8 @@ impl Default for ServerOptions {
             port: 5432,
             tls_cert_path: None,
             tls_key_path: None,
+            max_connections: 1000,
+            query_timeout: Some(std::time::Duration::from_secs(30)),
         }
     }
 }
@@ -85,7 +100,7 @@ pub async fn serve(
     let auth_manager = Arc::new(AuthManager::new());
 
     // Create the handler factory with authentication
-    let factory = Arc::new(HandlerFactory::new(session_context, auth_manager));
+    let factory = Arc::new(HandlerFactory::new(session_context, auth_manager, opts.query_timeout));
 
     serve_with_handlers(factory, opts).await
 }
@@ -126,17 +141,31 @@ pub async fn serve_with_handlers(
         info!("Listening on {server_addr} (unencrypted)");
     }
 
+    // Connection limiter to prevent resource exhaustion
+    let connection_semaphore = Arc::new(Semaphore::new(opts.max_connections));
+
     // Accept incoming connections
     loop {
         match listener.accept().await {
-            Ok((socket, _addr)) => {
+            Ok((socket, addr)) => {
                 let factory_ref = handlers.clone();
                 let tls_acceptor_ref = tls_acceptor.clone();
+                let semaphore_ref = connection_semaphore.clone();
 
                 tokio::spawn(async move {
+                    // Acquire connection permit to limit concurrency
+                    let _permit = match semaphore_ref.try_acquire() {
+                        Ok(permit) => permit,
+                        Err(_) => {
+                            warn!("Connection rejected from {addr}: max connections reached");
+                            return;
+                        }
+                    };
+
                     if let Err(e) = process_socket(socket, tls_acceptor_ref, factory_ref).await {
-                        warn!("Error processing socket: {e}");
+                        warn!("Error processing socket from {addr}: {e}");
                     }
+                    // Permit is automatically released when _permit is dropped
                 });
             }
             Err(e) => {