feat: update pgwire to 0.39

Author: sunng87

Cargo.lock

@@ -19,7 +19,7 @@ bytes = "1.11.1"
 chrono = { version = "0.4", features = ["std"] }
 datafusion = { version = "53" }
 futures = "0.3"
-pgwire = { version = "0.38", default-features = false }
+pgwire = { version = "0.39", default-features = false }
 postgres-types = "0.2"
 rust_decimal = { version = "1.41", features = ["db-postgres"] }
 tokio = { version = "1", default-features = false }

Cargo.toml

@@ -11,11 +11,12 @@ use datafusion::sql::sqlparser;
 use log::info;
 use pgwire::api::auth::noop::NoopStartupHandler;
 use pgwire::api::auth::StartupHandler;
+use pgwire::api::cancel::{CancelHandler, DefaultCancelHandler};
 use pgwire::api::portal::{Format, Portal};
 use pgwire::api::query::{ExtendedQueryHandler, SimpleQueryHandler};
 use pgwire::api::results::{FieldInfo, Response, Tag};
 use pgwire::api::stmt::QueryParser;
-use pgwire::api::{ClientInfo, ErrorHandler, PgWireServerHandlers, Type};
+use pgwire::api::{ClientInfo, ConnectionManager, ErrorHandler, PgWireServerHandlers, Type};
 use pgwire::error::{PgWireError, PgWireResult};
 use pgwire::messages::PgWireBackendMessage;
 use pgwire::types::format::FormatOptions;
@@ -29,19 +30,34 @@ use arrow_pg::datatypes::{arrow_schema_to_pg_fields, into_pg_type};
 use datafusion_pg_catalog::sql::PostgresCompatibilityParser;
 
 /// Simple startup handler that does no authentication
-pub struct SimpleStartupHandler;
+pub struct SimpleStartupHandler {
+    connection_manager: Arc<ConnectionManager>,
+}
 
 #[async_trait::async_trait]
-impl NoopStartupHandler for SimpleStartupHandler {}
+impl NoopStartupHandler for SimpleStartupHandler {
+    fn connection_manager(&self) -> Option<Arc<ConnectionManager>> {
+        Some(self.connection_manager.clone())
+    }
+}
 
 pub struct HandlerFactory {
     pub session_service: Arc<DfSessionService>,
+    cancel_handler: Arc<DefaultCancelHandler>,
+    startup_handler: Arc<SimpleStartupHandler>,
 }
 
 impl HandlerFactory {
     pub fn new(session_context: Arc<SessionContext>) -> Self {
         let session_service = Arc::new(DfSessionService::new(session_context));
-        HandlerFactory { session_service }
+        let connection_manager = Arc::new(ConnectionManager::new());
+        HandlerFactory {
+            session_service,
+            cancel_handler: Arc::new(DefaultCancelHandler::new(connection_manager.clone())),
+            startup_handler: Arc::new(SimpleStartupHandler {
+                connection_manager: connection_manager.clone(),
+            }),
+        }
     }
 
     pub fn new_with_hooks(
@@ -52,7 +68,14 @@ impl HandlerFactory {
             session_context,
             query_hooks,
         ));
-        HandlerFactory { session_service }
+        let connection_manager = Arc::new(ConnectionManager::new());
+        HandlerFactory {
+            session_service,
+            cancel_handler: Arc::new(DefaultCancelHandler::new(connection_manager.clone())),
+            startup_handler: Arc::new(SimpleStartupHandler {
+                connection_manager: connection_manager.clone(),
+            }),
+        }
     }
 }
 
@@ -66,12 +89,16 @@ impl PgWireServerHandlers for HandlerFactory {
     }
 
     fn startup_handler(&self) -> Arc<impl StartupHandler> {
-        Arc::new(SimpleStartupHandler)
+        self.startup_handler.clone()
     }
 
     fn error_handler(&self) -> Arc<impl ErrorHandler> {
         Arc::new(LoggingErrorHandler)
     }
+
+    fn cancel_handler(&self) -> Arc<impl CancelHandler> {
+        self.cancel_handler.clone()
+    }
 }
 
 struct LoggingErrorHandler;

datafusion-postgres/src/handlers.rs

@@ -4,7 +4,10 @@ use datafusion::prelude::{SessionConfig, SessionContext};
 use datafusion_pg_catalog::pg_catalog::setup_pg_catalog;
 use futures::Sink;
 use pgwire::{
-    api::{ClientInfo, ClientPortalStore, PgWireConnectionState, METADATA_USER},
+    api::{
+        store::MemPortalStore, ClientInfo, ClientPortalStore, PgWireConnectionState,
+        SessionExtensions, METADATA_USER,
+    },
     messages::{
         response::TransactionStatus, startup::SecretKey, PgWireBackendMessage, ProtocolVersion,
     },
@@ -29,8 +32,9 @@ pub fn setup_handlers() -> DfSessionService {
 #[derive(Debug, Default)]
 pub struct MockClient {
     metadata: HashMap<String, String>,
-    portal_store: HashMap<String, String>,
+    portal_store: MemPortalStore<String>,
     pub sent_messages: Vec<PgWireBackendMessage>,
+    session_extensions: SessionExtensions,
 }
 
 impl MockClient {
@@ -40,8 +44,9 @@ impl MockClient {
 
         MockClient {
             metadata,
-            portal_store: HashMap::default(),
+            portal_store: MemPortalStore::new(),
             sent_messages: Vec::new(),
+            session_extensions: SessionExtensions::new(),
         }
     }
 
@@ -98,10 +103,14 @@ impl ClientInfo for MockClient {
     fn sni_server_name(&self) -> Option<&str> {
         None
     }
+
+    fn session_extensions(&self) -> &pgwire::api::SessionExtensions {
+        &self.session_extensions
+    }
 }
 
 impl ClientPortalStore for MockClient {
-    type PortalStore = HashMap<String, String>;
+    type PortalStore = MemPortalStore<String>;
     fn portal_store(&self) -> &Self::PortalStore {
         &self.portal_store
     }