From ab33ca232f09c36c5883c7acea7c97c8a2311bc1 Mon Sep 17 00:00:00 2001 From: Erik Johnston Date: Thu, 2 Jul 2026 15:35:36 +0000 Subject: [PATCH] Apply per-connection session setup in the Rust connection pool Move the equivalent of the engine's `on_new_connection` into the pool's `ConnectionManager::create`, so every pooled connection is configured once when it is opened rather than on each checkout. A new `SessionConfig` (carried by the manager) drives it: - the default isolation level is set to REPEATABLE READ, matching the engine's `default_isolation_level`, so a plain BEGIN behaves as it did under psycopg2; - `synchronous_commit` is turned off when configured; - `statement_timeout` is set when configured. `bytea_output` is deliberately not set: tokio-postgres uses the binary protocol for prepared statements, so that text-format GUC is irrelevant to how we decode bytea. `create_pool` keeps its signature (default session); a new `create_pool_with_session` and two new keyword-only `ConnectionPool(...)` arguments (`synchronous_commit`, `statement_timeout_ms`) expose the settings. Co-Authored-By: Claude Opus 4.8 (1M context) Claude-Session: https://claude.ai/code/session_01W3G4M92AmwSSZCbmtMJU3d --- rust/src/database/postgres/pool.rs | 140 ++++++++++++++++++++++++++--- 1 file changed, 126 insertions(+), 14 deletions(-) diff --git a/rust/src/database/postgres/pool.rs b/rust/src/database/postgres/pool.rs index 6a21246f3b..6631d78361 100644 --- a/rust/src/database/postgres/pool.rs +++ b/rust/src/database/postgres/pool.rs @@ -3,9 +3,9 @@ //! //! We use the generic `deadpool::managed` pool with our own [`ConnectionManager`] //! rather than the `deadpool-postgres` crate, so that creating a connection -//! reuses the same DSN handling (libpq's default host, see -//! [`super::fixup_default_host`]) and connection-task spawning as -//! [`super::connect`]. +//! reuses our own DSN handling (libpq's default host, see +//! [`super::fixup_default_host`]) and drives the connection task on the shared +//! runtime. //! //! The pooled item is a plain [`tokio_postgres::Client`]. Rust-native code takes //! one from the pool and uses it with the standard `tokio_postgres` async query @@ -25,18 +25,68 @@ use crate::database::postgres::fixup_default_host; use crate::database::postgres::helpers::BlockingPostgres; use crate::database::runtime::runtime; +/// Per-connection session settings applied once when a pooled connection is +/// opened (the native equivalent of the engine's `on_new_connection`). +/// +/// Note we deliberately do *not* set `bytea_output`: unlike psycopg2, +/// [`tokio_postgres`] talks the binary protocol for prepared statements, so the +/// text-format `bytea_output` GUC is irrelevant to how we decode `bytea`. +#[derive(Clone)] +pub struct SessionConfig { + /// When false, `synchronous_commit` is turned off for the session (don't + /// wait for the server to fsync before returning from a commit). + pub synchronous_commit: bool, + /// If set, `statement_timeout` (in milliseconds) — statements running longer + /// are aborted and turned into errors. + pub statement_timeout_ms: Option, +} + +impl Default for SessionConfig { + fn default() -> Self { + // Match the engine's defaults: synchronous commit on, no statement + // timeout unless configured. + Self { + synchronous_commit: true, + statement_timeout_ms: None, + } + } +} + +impl SessionConfig { + /// The `SET` statements to run on a freshly-opened connection. + fn setup_sql(&self) -> String { + // Match the engine's `default_isolation_level` (REPEATABLE READ) as the + // session default, so a plain `BEGIN` gets the same isolation psycopg2 + // gave us. + let mut sql = String::from( + "SET SESSION CHARACTERISTICS AS TRANSACTION ISOLATION LEVEL REPEATABLE READ;", + ); + if !self.synchronous_commit { + sql.push_str(" SET synchronous_commit TO OFF;"); + } + if let Some(ms) = self.statement_timeout_ms { + // `ms` is an integer we control (from config), so inlining is safe. + sql.push_str(&format!(" SET statement_timeout TO {ms};")); + } + sql + } +} + /// Creates and recycles [`tokio_postgres`] connections for a [`ConnectionPool`]. pub struct ConnectionManager { /// The resolved connection config, parsed once from the DSN (with libpq's /// default host filled in if the DSN omitted one). config: Config, + /// Session settings applied to each connection when it is opened. + session: SessionConfig, } impl ConnectionManager { - /// Build a manager from a libpq-style DSN. - pub fn from_dsn(dsn: &str) -> Result { + /// Build a manager from a libpq-style DSN and session settings. + pub fn from_dsn(dsn: &str, session: SessionConfig) -> Result { Ok(Self { config: fixup_default_host(dsn)?, + session, }) } } @@ -46,10 +96,10 @@ impl Manager for ConnectionManager { type Error = tokio_postgres::Error; async fn create(&self) -> Result { - // As in `super::connect`: establish the connection, then drive its - // long-lived connection task (which pumps the socket) on the shared - // runtime. The task ends on its own when the `Client` is dropped, i.e. - // when the pool discards this connection. + // Establish the connection, then drive its long-lived connection task + // (which pumps the socket) on the shared runtime. The task ends on its + // own when the `Client` is dropped, i.e. when the pool discards this + // connection. let (client, connection) = self.config.connect(NoTls).await?; runtime().spawn(async move { @@ -58,6 +108,10 @@ impl Manager for ConnectionManager { } }); + // Apply per-connection session setup once, up front (the native + // equivalent of the engine's `on_new_connection`). + client.batch_execute(&self.session.setup_sql()).await?; + Ok(client) } @@ -81,9 +135,18 @@ pub type ConnectionPool = Pool; pub type PooledConnection = Object; /// Build a [`ConnectionPool`] from a libpq-style DSN, capped at `max_size` -/// connections. +/// connections, using the default [`SessionConfig`]. pub fn create_pool(dsn: &str, max_size: usize) -> Result { - let manager = ConnectionManager::from_dsn(dsn)?; + create_pool_with_session(dsn, max_size, SessionConfig::default()) +} + +/// Build a [`ConnectionPool`] with explicit per-connection [`SessionConfig`]. +pub fn create_pool_with_session( + dsn: &str, + max_size: usize, + session: SessionConfig, +) -> Result { + let manager = ConnectionManager::from_dsn(dsn, session)?; Ok(Pool::builder(manager).max_size(max_size).build()?) } @@ -104,10 +167,22 @@ impl PyConnectionPool { /// /// This only parses the DSN; connections are opened lazily on the first /// (and each subsequent) [`connect`](Self::connect) that needs a new one. + /// Each connection gets the same per-connection session setup: the default + /// (REPEATABLE READ) isolation level, plus `synchronous_commit` / + /// `statement_timeout` if configured here. #[new] - #[pyo3(signature = (dsn, max_size = 10))] - fn new(dsn: &str, max_size: usize) -> PyResult { - let pool = create_pool(dsn, max_size).map_err(|e| { + #[pyo3(signature = (dsn, max_size = 10, *, synchronous_commit = true, statement_timeout_ms = None))] + fn new( + dsn: &str, + max_size: usize, + synchronous_commit: bool, + statement_timeout_ms: Option, + ) -> PyResult { + let session = SessionConfig { + synchronous_commit, + statement_timeout_ms, + }; + let pool = create_pool_with_session(dsn, max_size, session).map_err(|e| { PyRuntimeError::new_err(format!("failed to build connection pool: {e}")) })?; Ok(Self { pool }) @@ -221,4 +296,41 @@ mod tests { ); }); } + + #[test] + fn create_applies_session_setup_to_each_connection() { + let Some(dsn) = test_dsn() else { + eprintln!("skipping: set SYNAPSE_TEST_POSTGRES_DSN to run"); + return; + }; + + let session = SessionConfig { + synchronous_commit: false, + statement_timeout_ms: Some(1234), + }; + let pool = create_pool_with_session(&dsn, 1, session).unwrap(); + + runtime().block_on(async { + let client = pool.get().await.unwrap(); + + // The session default isolation matches the engine's default, and + // the configured `synchronous_commit` / `statement_timeout` are set. + let show = |name: &'static str| { + let client = &client; + async move { + client + .query_one(&format!("SHOW {name}"), &[]) + .await + .unwrap() + .get::<_, String>(0) + } + }; + assert_eq!( + show("default_transaction_isolation").await, + "repeatable read" + ); + assert_eq!(show("synchronous_commit").await, "off"); + assert_eq!(show("statement_timeout").await, "1234ms"); + }); + } }