build: bump pgwire to 0.39 and migrate to SessionExtensions

pgwire 0.39.0 added typed per-connection state via `SessionExtensions`,
which was the upstream change `TODO.md` was waiting for. Migrate
`ConnectionState` storage off the global `DashMap` keyed by
`{addr}_{pid}` and onto `client.session_extensions()`.

Mechanical changes:

- `session.rs`: drop the `CONNECTIONS` global, the `_conn_id`
  metadata key, `register_connection`, `get_connection`, and
  `remove_connections_for_addr`. Cancel registry stays — a
  `CancelRequest` arrives on a *different* TCP connection that
  doesn't have access to the originating connection's
  `SessionExtensions`, so a process-wide map is still required.
  Replaced its `dashmap::DashMap` with `Mutex<HashMap>` since the
  registry is small and only touched on connect/disconnect/cancel.
- `cancel.rs`: `lookup_cancel` now returns the cloned `(client,
  active_query_id)` pair directly, so the registry mutex is
  released before any await and the previous "snapshot then drop
  the ref guard" dance is gone.
- `startup.rs::establish_session`: replaces `register_connection`
  with `client.session_extensions().insert(ConnectionState { ... })`.
  No more `_conn_id` formatting or metadata insertion.
- `query_simple.rs` / `query_extended.rs`: replace
  `session::get_connection(&conn_id)` with
  `client.session_extensions().get::<ConnectionState>()`. The
  returned `Arc<ConnectionState>` is `Send + Sync + 'static`, so we
  drop the per-handler `Arc::clone` dance and the
  `drop(conn_state)` before-await workaround.
- `main.rs`: drop `ConnectionCleanup` (the per-task drop guard that
  called `remove_connections_for_addr`). Cleanup happens via
  `Drop for ConnectionState` when pgwire tears down the connection's
  `SessionExtensions`.
- `Cargo.toml`: drop `dashmap` direct dependency (still pulled in
  transitively by trino-rust-client).
- `SimpleQueryHandler` no longer requires `ClientPortalStore` in
  0.39 — drop the trait bound.

Also bumps tokio (1.52.0 -> 1.52.1) and rustls-pki-types (1.14.0 ->
1.14.1) to their latest patch revisions.

`cargo build --all-targets`, `cargo test --lib` (191 passed),
`cargo test --bin` (5 passed), `cargo clippy --all-targets -- -D
warnings`, `cargo fmt --check`, and `reuse lint` (59/59) all pass
post-migration.

Net: -125 lines (244 deletions, 119 insertions).

Author: lfrancke

Cargo.lock

@@ -10,7 +10,7 @@ repository = "https://github.com/stackabletech/postgresql-trino-gateway"
 keywords = ["postgresql", "trino", "gateway", "powerbi"]
 
 [dependencies]
-pgwire = "0.38"
+pgwire = "0.39"
 tokio = { version = "1", features = ["full"] }
 async-trait = "0.1"
 futures = "0.3"
@@ -21,7 +21,6 @@ anyhow = "1"
 serde_json = "1"
 trino-rust-client = "0.9"
 async-stream = "0.3"
-dashmap = "6"
 sqlparser = { version = "0.61", features = ["visitor"] }
 rustls-pemfile = "2"
 rustls-pki-types = "1"

Cargo.toml

@@ -26,26 +26,22 @@ pub struct GatewayCancelHandler;
 #[async_trait]
 impl CancelHandler for GatewayCancelHandler {
     async fn on_cancel_request(&self, req: CancelRequest) {
-        let entry = match session::lookup_cancel(req.pid, &req.secret_key) {
-            Some(e) => e,
-            None => {
-                tracing::debug!(
-                    pid = req.pid,
-                    "CancelRequest for unknown (pid, secret_key) — ignored"
-                );
-                return;
-            }
+        let Some((trino_client, active_query_id)) =
+            session::lookup_cancel(req.pid, &req.secret_key)
+        else {
+            tracing::debug!(
+                pid = req.pid,
+                "CancelRequest for unknown (pid, secret_key) — ignored"
+            );
+            return;
         };
 
-        // Snapshot the query id and trino-client handle, then drop the
-        // DashMap ref guard before the async cancel call — guards aren't
-        // Send and can't be held across awaits.
-        let trino_client = std::sync::Arc::clone(&entry.trino_client);
-        let query_id = match entry.active_query_id.lock() {
+        // `lookup_cancel` already cloned out the bits we need; the registry
+        // mutex was released inside it, so we're free to await below.
+        let query_id = match active_query_id.lock() {
             Ok(g) => g.clone(),
             Err(p) => p.into_inner().clone(),
         };
-        drop(entry);
 
         let Some(qid) = query_id else {
             tracing::debug!(

src/cancel.rs

@@ -1,6 +1,5 @@
 // SPDX-FileCopyrightText: 2026 Stackable GmbH
 // SPDX-License-Identifier: OSL-3.0
-use std::net::SocketAddr;
 use std::sync::Arc;
 use std::time::Duration;
 
@@ -15,7 +14,6 @@ use postgresql_trino_gateway::handler::GatewayHandlerFactory;
 use postgresql_trino_gateway::policy;
 use postgresql_trino_gateway::query_extended::GatewayExtendedQueryHandler;
 use postgresql_trino_gateway::query_simple::GatewayQueryHandler;
-use postgresql_trino_gateway::session;
 use postgresql_trino_gateway::startup::GatewayStartupHandler;
 use postgresql_trino_gateway::tls;
 
@@ -31,16 +29,6 @@ const SHUTDOWN_DRAIN_TIMEOUT_ENV: &str = "GATEWAY_SHUTDOWN_DRAIN_TIMEOUT_SECS";
 /// unwinding before letting the process exit.
 const ABORT_UNWIND_GRACE: Duration = Duration::from_secs(1);
 
-/// Drop guard that removes per-connection state when the connection task
-/// ends, including on panic or cancellation. Defined at module scope so the
-/// accept loop body stays readable.
-struct ConnectionCleanup(SocketAddr);
-impl Drop for ConnectionCleanup {
-    fn drop(&mut self) {
-        session::remove_connections_for_addr(self.0);
-    }
-}
-
 #[tokio::main]
 async fn main() -> anyhow::Result<()> {
     tracing_subscriber::fmt::init();
@@ -153,9 +141,8 @@ async fn main() -> anyhow::Result<()> {
                 let tls = tls_acceptor.clone();
                 tasks.spawn(async move {
                     let _permit = permit; // released on task exit
-                    let _guard = ConnectionCleanup(peer_addr);
                     if let Err(e) = pgwire::tokio::process_socket(socket, tls, factory).await {
-                        tracing::error!(error = %e, "connection error");
+                        tracing::error!(addr = %peer_addr, error = %e, "connection error");
                     }
                 });
             }

src/main.rs

@@ -15,7 +15,7 @@ use pgwire::error::{PgWireError, PgWireResult};
 use pgwire::messages::PgWireBackendMessage;
 
 use crate::query_pipeline::process_query;
-use crate::session::{self, CachedPortalResponse, MAX_CACHED_PORTALS, PortalCache};
+use crate::session::{CachedPortalResponse, ConnectionState, MAX_CACHED_PORTALS, PortalCache};
 
 /// Handles the extended query protocol (Parse/Bind/Describe/Execute).
 ///
@@ -86,15 +86,11 @@ impl ExtendedQueryHandler for GatewayExtendedQueryHandler {
         PgWireError: From<<C as Sink<PgWireBackendMessage>>::Error>,
     {
         let query = &portal.statement.statement;
+        tracing::debug!(query, "Extended query execute");
 
-        let conn_id = client
-            .metadata()
-            .get(session::CONNECTION_ID_KEY)
-            .ok_or_else(|| PgWireError::ApiError("No connection ID in metadata".into()))?
-            .clone();
-        tracing::debug!(conn_id = %conn_id, query, "Extended query execute");
-
-        let conn_state = session::get_connection(&conn_id)
+        let conn_state = client
+            .session_extensions()
+            .get::<ConnectionState>()
             .ok_or_else(|| PgWireError::ApiError("Connection state not found".into()))?;
 
         // Take any response stashed by do_describe_portal. If the cached
@@ -104,39 +100,36 @@ impl ExtendedQueryHandler for GatewayExtendedQueryHandler {
         let cached_entry = take_cached(&conn_state.portals, &portal.name)?;
         if let Some(entry) = cached_entry {
             if entry.query == *query {
-                tracing::trace!(conn_id = %conn_id, portal = %portal.name, "Extended query execute: served from describe cache");
+                tracing::trace!(portal = %portal.name, "Extended query execute: served from describe cache");
                 return Ok(entry.response);
             }
             tracing::debug!(
-                conn_id = %conn_id,
                 portal = %portal.name,
                 "Discarding stale describe-cache entry — portal was re-bound"
             );
         }
 
-        let trino_client = Arc::clone(&conn_state.trino_client);
-        let config = Arc::clone(&conn_state.config);
-        let active_query_id = Arc::clone(&conn_state.active_query_id);
-        drop(conn_state);
-
-        let responses =
-            process_query(query, &trino_client, &config, Some(&active_query_id)).await?;
+        let responses = process_query(
+            query,
+            &conn_state.trino_client,
+            &conn_state.config,
+            Some(&conn_state.active_query_id),
+        )
+        .await?;
         let response = responses
             .into_iter()
             .next()
             .ok_or_else(|| PgWireError::ApiError("Empty pipeline response".into()))?;
         match &response {
             Response::Query(qr) => {
                 tracing::trace!(
-                    conn_id = %conn_id,
                     portal = %portal.name,
                     columns = qr.row_schema.len(),
                     "Extended query execute: query response (no cache)"
                 );
             }
             Response::Execution(_tag) => {
                 tracing::trace!(
-                    conn_id = %conn_id,
                     portal = %portal.name,
                     "Extended query execute: execution response (no cache)"
                 );
@@ -185,25 +178,20 @@ impl ExtendedQueryHandler for GatewayExtendedQueryHandler {
         // Run the pipeline once, return the schema, and stash the response so
         // do_query can serve Execute without re-running.
         let query = &portal.statement.statement;
+        tracing::debug!(query, "Extended query describe portal");
 
-        let conn_id = client
-            .metadata()
-            .get(session::CONNECTION_ID_KEY)
-            .ok_or_else(|| PgWireError::ApiError("No connection ID in metadata".into()))?
-            .clone();
-        tracing::debug!(conn_id = %conn_id, query, "Extended query describe portal");
-        let conn_state = session::get_connection(&conn_id)
+        let conn_state = client
+            .session_extensions()
+            .get::<ConnectionState>()
             .ok_or_else(|| PgWireError::ApiError("Connection state not found".into()))?;
-        let trino_client = Arc::clone(&conn_state.trino_client);
-        let config = Arc::clone(&conn_state.config);
-        // Clone the Arcs so we can use them after the .await below;
-        // conn_state is a DashMap ref-guard that can't be held across awaits.
-        let portals = Arc::clone(&conn_state.portals);
-        let active_query_id = Arc::clone(&conn_state.active_query_id);
-        drop(conn_state);
 
-        let responses =
-            process_query(query, &trino_client, &config, Some(&active_query_id)).await?;
+        let responses = process_query(
+            query,
+            &conn_state.trino_client,
+            &conn_state.config,
+            Some(&conn_state.active_query_id),
+        )
+        .await?;
         let response = responses
             .into_iter()
             .next()
@@ -228,7 +216,7 @@ impl ExtendedQueryHandler for GatewayExtendedQueryHandler {
         // Trino's `DELETE /v1/statement/{queryId}` so abandoned queries
         // release Trino resources promptly.
         insert_cached(
-            &portals,
+            &conn_state.portals,
             portal.name.clone(),
             CachedPortalResponse {
                 query: query.clone(),

src/query_extended.rs

@@ -1,19 +1,17 @@
 // SPDX-FileCopyrightText: 2026 Stackable GmbH
 // SPDX-License-Identifier: OSL-3.0
 use std::fmt::Debug;
-use std::sync::Arc;
 
 use async_trait::async_trait;
 use futures::Sink;
+use pgwire::api::ClientInfo;
+use pgwire::api::query::SimpleQueryHandler;
 use pgwire::api::results::Response;
-use pgwire::api::{ClientInfo, ClientPortalStore};
 use pgwire::error::{PgWireError, PgWireResult};
 use pgwire::messages::PgWireBackendMessage;
 
-use pgwire::api::query::SimpleQueryHandler;
-
 use crate::query_pipeline::process_query;
-use crate::session;
+use crate::session::ConnectionState;
 
 #[derive(Debug)]
 pub struct GatewayQueryHandler;
@@ -22,37 +20,29 @@ pub struct GatewayQueryHandler;
 impl SimpleQueryHandler for GatewayQueryHandler {
     async fn do_query<C>(&self, client: &mut C, query: &str) -> PgWireResult<Vec<Response>>
     where
-        C: ClientInfo + ClientPortalStore + Sink<PgWireBackendMessage> + Unpin + Send + Sync,
+        C: ClientInfo + Sink<PgWireBackendMessage> + Unpin + Send + Sync,
         C::Error: Debug,
         PgWireError: From<<C as Sink<PgWireBackendMessage>>::Error>,
     {
         tracing::debug!(query, "Simple query received");
 
-        let conn_id = client
-            .metadata()
-            .get(session::CONNECTION_ID_KEY)
-            .ok_or_else(|| PgWireError::ApiError("No connection ID in metadata".into()))?
-            .clone();
-        let conn_state = session::get_connection(&conn_id)
+        let conn_state = client
+            .session_extensions()
+            .get::<ConnectionState>()
             .ok_or_else(|| PgWireError::ApiError("Connection state not found".into()))?;
-        let trino_client = Arc::clone(&conn_state.trino_client);
-        let config = Arc::clone(&conn_state.config);
-        let active_query_id = Arc::clone(&conn_state.active_query_id);
-        // `conn_state` is a `dashmap::Ref` guard holding a read-lock on its
-        // shard; the guard is not `Send`, so it cannot be held across the
-        // `process_query(...).await` below. Drop it explicitly after we've
-        // cloned out the Arcs we need; otherwise the borrow-checker rejects
-        // the await.
-        drop(conn_state);
 
-        let result = process_query(query, &trino_client, &config, Some(&active_query_id)).await;
+        let result = process_query(
+            query,
+            &conn_state.trino_client,
+            &conn_state.config,
+            Some(&conn_state.active_query_id),
+        )
+        .await;
         match &result {
-            Ok(responses) => tracing::trace!(
-                conn_id = %conn_id,
-                response_count = responses.len(),
-                "Simple query processed"
-            ),
-            Err(e) => tracing::debug!(conn_id = %conn_id, error = %e, "Simple query failed"),
+            Ok(responses) => {
+                tracing::trace!(response_count = responses.len(), "Simple query processed")
+            }
+            Err(e) => tracing::debug!(error = %e, "Simple query failed"),
         }
         result
     }