Make the Postgres Connection shim pool-backed

The Python-facing `Connection`/`Cursor` shim now always wraps a connection
checked out of the `deadpool` pool. `ConnInner` holds a `PooledConnection` and
disposes of it correctly when the last reference (the `Connection` and all its
cursors) goes away, via a `Drop` impl plus `release`/`discard` helpers:

  - a clean connection is returned to the pool for reuse;
  - a connection is *discarded* (detached with `Object::take`, which also
    shrinks the pool) whenever reuse would be unsafe — a failed COMMIT/ROLLBACK,
    a poisoned mutex, or being dropped with a transaction still open (which
    can't be rolled back synchronously from `Drop`, so the socket close makes
    the server do it).

A plain query error still does *not* throw the connection away — it stays
open+aborted for the driver to `rollback()`, exactly as psycopg2 behaves.

The pool is the only way to obtain a connection: the standalone `connect(dsn)`
free function is replaced by a Python-facing `ConnectionPool` class (Rust
`PyConnectionPool`, exposed as `postgres.ConnectionPool`). Build it once from a
DSN, then check connections out with `pool.connect()`. Checkout failures map
onto the DBAPI2 hierarchy — a backend connect error reuses the query-error
mapping (and its `pgcode`), while a timeout / closed pool becomes
`OperationalError` — so `connect()` behaves like psycopg2's for callers.

Adds live-Postgres tests (gated on SYNAPSE_TEST_POSTGRES_DSN) asserting which
connections end up back in the pool, and drives the Python test suite through a
pool in `setUp`.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01W3G4M92AmwSSZCbmtMJU3d

rust/src/database/postgres/connection.rs

@@ -39,21 +39,34 @@
 //! `commit`/`rollback` end the transaction and clear the flag; with no
 //! transaction open they are no-ops, just like psycopg2.
 //!
-//! ## Dropping the `Client` on error
+//! ## Returning vs discarding the connection
+//!
+//! Every [`Connection`] wraps a [`PooledConnection`] checked out of the
+//! [`super::pool`]. Dropping it normally *returns it to the pool* for reuse.
+//! Where that reuse would be unsafe we instead **discard** it: the connection is
+//! detached from the pool with [`Object::take`] and dropped, which both closes
+//! the socket and shrinks the pool so the bad connection is never handed out
+//! again.
 //!
 //! A *query* error (bad SQL, a constraint violation, an integer out of range,
 //! …) leaves the connection open with its transaction in the aborted state,
 //! exactly as psycopg2 does: the error propagates to Python and the driver is
 //! expected to `rollback()`. We do **not** throw the connection away for these.
 //!
-//! The transaction-control statements are different. If `COMMIT` or `ROLLBACK`
-//! itself fails we no longer know what state the server-side session is in, so
-//! we drop the `Client` (closing the socket) rather than hand a possibly-broken
-//! connection back for reuse. Likewise, [`Connection::close`] drops the client;
-//! the server rolls back any transaction left open when the socket closes.
+//! Three situations do force a discard, because the session state is unknown or
+//! unclean and must not reach the next caller:
+//!  - a failed `COMMIT`/`ROLLBACK` — we no longer know the session state;
+//!  - a poisoned mutex — a panic happened mid-operation;
+//!  - the connection being dropped with a transaction still open — it can't be
+//!    rolled back synchronously from `Drop`, so the server does it for us when
+//!    the socket closes.
+//!
+//! [`Connection::close`], by contrast, returns a clean connection to the pool
+//! (discarding it only if a transaction was left open).
 
 use std::sync::{Arc, Mutex, MutexGuard, TryLockError};
 
+use deadpool::managed::Object;
 use futures::future::try_join_all;
 use pyo3::{
     exceptions::PyRuntimeError,
@@ -63,7 +76,8 @@ use pyo3::{
 use tokio_postgres::Client;
 
 use crate::database::postgres::{
-    cursor_state::CursorQueryState, helpers::BlockingPostgresResult, value::PgValue,
+    cursor_state::CursorQueryState, helpers::BlockingPostgresResult, pool::PooledConnection,
+    value::PgValue,
 };
 
 /// `try_lock` a mutex that is single-threaded by contract, mapping its two
@@ -94,10 +108,10 @@ fn try_lock_or_reset<'a, T>(
 
 /// A single Postgres connection exposed to Python.
 ///
-/// Owns the [`tokio_postgres::Client`] for its whole life and is the authority
-/// on transaction state. The `Arc<Mutex<...>>` lets cursors hold a cheap clone
-/// (so they can reach the client to start a query) while keeping all access to
-/// the client serialised.
+/// Wraps a connection checked out of the pool for its whole life and is the
+/// authority on transaction state. The `Arc<Mutex<...>>` lets cursors hold a
+/// cheap clone (so they can reach the client to start a query) while keeping all
+/// access to the client serialised.
 #[pyclass(frozen, skip_from_py_object)]
 #[derive(Clone)]
 pub struct Connection {
@@ -106,9 +120,10 @@ pub struct Connection {
 
 /// The mutable guts of a [`Connection`], behind its mutex.
 struct ConnInner {
-    /// The driver client. `None` once the connection has been closed (or thrown
-    /// away after a transaction-control error); any further use is an error.
-    client: Option<Client>,
+    /// The pooled connection. `None` once the connection has been closed,
+    /// returned to the pool, or discarded after an error; any further use is an
+    /// error.
+    client: Option<PooledConnection>,
     /// Whether a transaction is currently open (a `BEGIN` has been issued and
     /// not yet matched by a `COMMIT`/`ROLLBACK`). Drives the lazy `BEGIN`.
     in_txn: bool,
@@ -118,12 +133,51 @@ struct ConnInner {
     autocommit: bool,
 }
 
+impl ConnInner {
+    /// Give up the connection cleanly.
+    ///
+    /// The connection is returned to the pool for reuse — **unless** a
+    /// transaction is still open, in which case it can't be rolled back from
+    /// here, so we discard it (detach and drop) rather than hand a
+    /// mid-transaction connection to the next caller.
+    fn release(&mut self) {
+        if let Some(conn) = self.client.take() {
+            if self.in_txn {
+                let _ = Object::take(conn); // detach + drop: not returned to the pool
+            }
+            // else: `conn` dropped here → returned to the pool for reuse.
+        }
+        self.in_txn = false;
+    }
+
+    /// Discard the connection: the session state is unknown or unclean, so it
+    /// must never be reused. It is detached from the pool (shrinking it).
+    fn discard(&mut self) {
+        if let Some(conn) = self.client.take() {
+            let _ = Object::take(conn);
+        }
+        self.in_txn = false;
+    }
+}
+
+impl Drop for ConnInner {
+    fn drop(&mut self) {
+        // Return the connection to the pool (or discard it) when the last
+        // reference — the `Connection` and all its cursors — goes away.
+        self.release();
+    }
+}
+
 impl Connection {
-    /// Wrap a freshly-established `Client` in a `Connection`.
-    pub fn new(client: Client) -> Self {
+    /// Wrap a connection checked out of the pool in a `Connection`.
+    ///
+    /// The connection is returned to the pool when this `Connection` (and every
+    /// cursor cloned from it) is dropped, unless it is discarded first (see the
+    /// module docs).
+    pub fn new(conn: PooledConnection) -> Self {
         Self {
             inner: Arc::new(Mutex::new(ConnInner {
-                client: Some(client),
+                client: Some(conn),
                 in_txn: false,
                 autocommit: false,
             })),
@@ -140,10 +194,9 @@ impl Connection {
     /// know the session state — and errors.
     fn lock(&self) -> PyResult<MutexGuard<'_, ConnInner>> {
         try_lock_or_reset(&self.inner, "connection", |inner| {
-            // On poison we no longer know the session state, so close the
-            // connection: drop the client and clear the transaction flag.
-            inner.client = None;
-            inner.in_txn = false;
+            // On poison we no longer know the session state, so discard the
+            // connection (never returning it to the pool).
+            inner.discard();
         })
     }
 
@@ -208,9 +261,9 @@ impl Connection {
                 Ok(())
             }
             Err(err) => {
-                // Unknown session state: drop the connection rather than reuse it.
-                guard.client = None;
-                guard.in_txn = false;
+                // Unknown session state: discard the connection rather than
+                // reuse it (or return it to the pool).
+                guard.discard();
                 Err(err)
             }
         }
@@ -222,7 +275,7 @@ impl Connection {
 fn client_ref(guard: &ConnInner) -> PyResult<&Client> {
     guard
         .client
-        .as_ref()
+        .as_deref()
         .ok_or_else(|| PyRuntimeError::new_err("connection already closed"))
 }
 
@@ -247,15 +300,14 @@ impl Connection {
         self.end_txn(py, "ROLLBACK")
     }
 
-    /// Close the connection, dropping the underlying client.
+    /// Close the connection, releasing the underlying client.
     ///
-    /// Dropping the client closes the socket; the server rolls back any
-    /// transaction that was still open. Idempotent: closing an
+    /// A standalone client's socket is closed; a pooled connection is returned
+    /// to the pool for reuse (or discarded if a transaction was left open, in
+    /// which case the server rolls it back). Idempotent: closing an
     /// already-closed connection is fine.
     fn close(&self) -> PyResult<()> {
-        let mut guard = self.lock()?;
-        guard.client = None;
-        guard.in_txn = false;
+        self.lock()?.release();
         Ok(())
     }
 
@@ -556,3 +608,112 @@ impl Cursor {
         Ok(())
     }
 }
+
+#[cfg(test)]
+mod tests {
+    //! These exercise the pool-backed `Connection` against a live Postgres, so
+    //! they only run when `SYNAPSE_TEST_POSTGRES_DSN` is set (e.g. to
+    //! `host=postgres user=postgres password=postgres dbname=postgres`);
+    //! otherwise they no-op. They assert *which* connections end up back in the
+    //! pool — the transaction/value logic itself is covered by the Python test
+    //! suite that drives these classes end to end.
+
+    use super::*;
+    use crate::database::postgres::pool::create_pool;
+    use crate::database::runtime::runtime;
+
+    fn test_dsn() -> Option<String> {
+        std::env::var("SYNAPSE_TEST_POSTGRES_DSN").ok()
+    }
+
+    /// A clean connection (its transaction committed) is returned to the pool
+    /// when the `Connection` is dropped.
+    #[test]
+    fn pooled_connection_returns_to_pool_after_commit() {
+        let Some(dsn) = test_dsn() else {
+            eprintln!("skipping: set SYNAPSE_TEST_POSTGRES_DSN to run");
+            return;
+        };
+
+        let pool = create_pool(&dsn, 1).unwrap();
+        let obj = runtime().block_on(async { pool.get().await.unwrap() });
+        assert_eq!(pool.status().size, 1);
+
+        Python::initialize();
+        Python::attach(move |py| {
+            let conn = Connection::new(obj);
+            let cursor = conn.cursor();
+            cursor.execute(py, "SELECT 1", None).unwrap();
+            conn.commit(py).unwrap();
+            // Dropping both references releases the pooled connection.
+            drop(cursor);
+            drop(conn);
+        });
+
+        // The (clean) connection went back to the pool rather than being torn
+        // down, so it's available for the next caller.
+        assert_eq!(pool.status().size, 1);
+        assert_eq!(pool.status().available, 1);
+    }
+
+    /// A connection dropped with a transaction still open can't be rolled back
+    /// from `Drop`, so it is discarded (detached from the pool) rather than
+    /// handed to the next caller mid-transaction.
+    #[test]
+    fn pooled_connection_discarded_when_dropped_mid_transaction() {
+        let Some(dsn) = test_dsn() else {
+            eprintln!("skipping: set SYNAPSE_TEST_POSTGRES_DSN to run");
+            return;
+        };
+
+        let pool = create_pool(&dsn, 1).unwrap();
+        let obj = runtime().block_on(async { pool.get().await.unwrap() });
+        assert_eq!(pool.status().size, 1);
+
+        Python::initialize();
+        Python::attach(move |py| {
+            let conn = Connection::new(obj);
+            let cursor = conn.cursor();
+            // Opens a transaction lazily (BEGIN) but never commits/rolls back.
+            cursor.execute(py, "SELECT 1", None).unwrap();
+            drop(cursor);
+            drop(conn);
+        });
+
+        // Detached: the pool shrank rather than accepting a mid-transaction
+        // connection back.
+        assert_eq!(pool.status().size, 0);
+        assert_eq!(pool.status().available, 0);
+    }
+
+    /// A plain query error does *not* poison the connection: after the caller
+    /// rolls back, the (now-clean) connection returns to the pool.
+    #[test]
+    fn pooled_connection_returns_to_pool_after_query_error_and_rollback() {
+        let Some(dsn) = test_dsn() else {
+            eprintln!("skipping: set SYNAPSE_TEST_POSTGRES_DSN to run");
+            return;
+        };
+
+        let pool = create_pool(&dsn, 1).unwrap();
+        let obj = runtime().block_on(async { pool.get().await.unwrap() });
+
+        Python::initialize();
+        Python::attach(move |py| {
+            let conn = Connection::new(obj);
+            let cursor = conn.cursor();
+            // A bad statement aborts the transaction but leaves the connection
+            // usable, exactly as psycopg2 does.
+            cursor
+                .execute(py, "SELECT * FROM does_not_exist", None)
+                .unwrap_err();
+            // The driver's job on failure: roll back, which clears the txn.
+            conn.rollback(py).unwrap();
+            drop(cursor);
+            drop(conn);
+        });
+
+        assert_eq!(pool.status().size, 1);
+        assert_eq!(pool.status().available, 1);
+    }
+}

rust/src/database/postgres/mod.rs

@@ -4,14 +4,9 @@
 //! shared multi-thread tokio runtime (see `super::runtime`).
 
 use anyhow::Error;
-use log::warn;
-use pyo3::exceptions::PyRuntimeError;
 use pyo3::prelude::*;
 use pyo3::types::PyModule;
 
-use crate::database::postgres::helpers::BlockingPostgresResult;
-use crate::database::runtime::runtime;
-
 mod connection;
 mod cursor_state;
 mod errors;
@@ -21,15 +16,15 @@ pub mod pool;
 pub(crate) mod query;
 mod value;
 
-/// Register the `postgres` submodule (the `Connection` / `Cursor` classes, the
-/// DBAPI2 exception hierarchy and the `connect` factory) under the parent
+/// Register the `postgres` submodule (the `ConnectionPool`, `Connection` and
+/// `Cursor` classes and the DBAPI2 exception hierarchy) under the parent
 /// `database` module.
 pub fn register_module(py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> {
     let child = PyModule::new(py, "postgres")?;
 
+    child.add_class::<pool::PyConnectionPool>()?;
     child.add_class::<connection::Connection>()?;
     child.add_class::<connection::Cursor>()?;
-    child.add_function(wrap_pyfunction!(connect, &child)?)?;
     errors::register_exceptions(py, &child)?;
 
     m.add_submodule(&child)?;
@@ -43,33 +38,6 @@ pub fn register_module(py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()>
     Ok(())
 }
 
-/// Open a new Postgres connection from a libpq-style DSN.
-///
-/// Blocks until the connection is established, then spawns the long-lived
-/// connection task (which drives the socket) onto the shared runtime and
-/// hands back a `Connection` wrapping the client.
-#[pyfunction]
-fn connect<'py>(py: Python<'py>, dsn: &str) -> PyResult<Bound<'py, connection::Connection>> {
-    let config = fixup_default_host(dsn)
-        .map_err(|e| PyRuntimeError::new_err(format!("Failed to prepare DSN: {e}")))?;
-
-    // TLS is not yet supported: unlike libpq (whose default is
-    // `sslmode=prefer`), we never negotiate TLS regardless of the DSN's
-    // sslmode. Supporting it is left to a follow-up.
-    let (client, connection) = config.connect(tokio_postgres::NoTls).block_on_result(py)?;
-
-    // Spawn the connection task on the runtime.
-    runtime().spawn(async move {
-        if let Err(e) = connection.await {
-            warn!("postgres connection error: {e}");
-        }
-    });
-
-    let conn = connection::Connection::new(client);
-
-    Bound::new(py, conn)
-}
-
 /// Fix up a DSN to ensure it has a host, using libpq's default host if
 /// necessary.
 ///

rust/src/database/postgres/pool.rs

@@ -13,11 +13,16 @@
 //! *same* pool, so both share a single set of connections rather than running
 //! two pools that could exhaust the server's connection limit between them.
 
-use deadpool::managed::{Manager, Metrics, Object, Pool, RecycleError, RecycleResult};
+use deadpool::managed::{Manager, Metrics, Object, Pool, PoolError, RecycleError, RecycleResult};
 use log::warn;
+use pyo3::exceptions::PyRuntimeError;
+use pyo3::prelude::*;
 use tokio_postgres::{Client, Config, NoTls};
 
+use crate::database::postgres::connection::Connection;
+use crate::database::postgres::errors::{pg_err_to_py, OperationalError};
 use crate::database::postgres::fixup_default_host;
+use crate::database::postgres::helpers::BlockingPostgres;
 use crate::database::runtime::runtime;
 
 /// Creates and recycles [`tokio_postgres`] connections for a [`ConnectionPool`].
@@ -82,6 +87,57 @@ pub fn create_pool(dsn: &str, max_size: usize) -> Result<ConnectionPool, anyhow:
     Ok(Pool::builder(manager).max_size(max_size).build()?)
 }
 
+/// The Python-facing connection pool.
+///
+/// This is the single entry point Python uses to obtain a [`Connection`]:
+/// build a pool from a DSN once, then check connections out of it with
+/// [`PyConnectionPool::connect`]. Each checkout hands back a [`Connection`]
+/// borrowed from the pool that returns itself for reuse when closed or dropped.
+#[pyclass(name = "ConnectionPool", frozen)]
+pub struct PyConnectionPool {
+    pool: ConnectionPool,
+}
+
+#[pymethods]
+impl PyConnectionPool {
+    /// Build a pool from a libpq-style DSN, capped at `max_size` connections.
+    ///
+    /// This only parses the DSN; connections are opened lazily on the first
+    /// (and each subsequent) [`connect`](Self::connect) that needs a new one.
+    #[new]
+    #[pyo3(signature = (dsn, max_size = 10))]
+    fn new(dsn: &str, max_size: usize) -> PyResult<Self> {
+        let pool = create_pool(dsn, max_size).map_err(|e| {
+            PyRuntimeError::new_err(format!("failed to build connection pool: {e}"))
+        })?;
+        Ok(Self { pool })
+    }
+
+    /// Check a connection out of the pool.
+    ///
+    /// Blocks (releasing the GIL) until a connection is available, opening a new
+    /// one if the pool is below `max_size`. A failure to establish the
+    /// connection surfaces through the same DBAPI2 exception hierarchy as a
+    /// query error, so callers can treat it like psycopg2's `connect`.
+    fn connect(&self, py: Python<'_>) -> PyResult<Connection> {
+        let conn = self.pool.get().block_on(py).map_err(pool_err_to_py)?;
+        Ok(Connection::new(conn))
+    }
+}
+
+/// Map a `deadpool` checkout failure onto the DBAPI2 exception hierarchy.
+fn pool_err_to_py(err: PoolError<tokio_postgres::Error>) -> PyErr {
+    match err {
+        // The backend failed to establish the connection: reuse the exact
+        // mapping (and `pgcode` tagging) a query error gets.
+        PoolError::Backend(e) => pg_err_to_py(&e),
+        // Timed out waiting for a slot, pool closed, no runtime, or a
+        // post-create hook failure: all connection-level problems, which
+        // Synapse treats as retryable operational errors.
+        other => OperationalError::new_err(format!("failed to acquire connection: {other}")),
+    }
+}
+
 #[cfg(test)]
 mod tests {
     //! These tests need a live Postgres, so they only run when

tests/synapse_rust/test_database_postgres.py

@@ -84,23 +84,31 @@ class PostgresConnectionTestCase(unittest.TestCase):
     """Tests for the Rust Postgres ``Connection`` / ``Cursor``."""
 
     def setUp(self) -> None:
-        self.conn = postgres.connect(_build_dsn())
+        # Connections are only ever obtained from a pool; check one out for the
+        # duration of the test.
+        self.pool = postgres.ConnectionPool(_build_dsn())
+        self.conn = self.pool.connect()
 
     def tearDown(self) -> None:
-        # Explicitly drop the connection to ensure that the underlying Rust
-        # object is dropped before the Python interpreter shuts down. Otherwise,
-        # the open connection will block us tearing down the test database.
+        # Explicitly drop the connection (returning it to the pool) and then the
+        # pool, so the underlying Rust objects are dropped before the Python
+        # interpreter shuts down. Otherwise the open connection would block us
+        # tearing down the test database.
         del self.conn
+        del self.pool
 
     # ------------------------------------------------------------------
-    # connect()
+    # ConnectionPool.connect()
     # ------------------------------------------------------------------
 
     def test_connect_bad_dsn_raises(self) -> None:
         # A syntactically valid but unconnectable DSN should raise one of our
-        # DBAPI2 errors rather than return a half-open connection.
+        # DBAPI2 errors rather than return a half-open connection. The pool
+        # parses the DSN eagerly but only dials on checkout, so the failure
+        # surfaces from connect().
+        pool = postgres.ConnectionPool("host=127.0.0.1 port=1 dbname=does_not_exist")
         with self.assertRaises(postgres.Error):
-            postgres.connect("host=127.0.0.1 port=1 dbname=does_not_exist")
+            pool.connect()
 
     # ------------------------------------------------------------------
     # execute() / fetch_one() / fetch_all()
@@ -783,12 +791,15 @@ class PostgresConnectionDrivenTestCase(unittest.TestCase):
     """
 
     def setUp(self) -> None:
-        self.conn = postgres.connect(_build_dsn())
+        self.pool = postgres.ConnectionPool(_build_dsn())
+        self.conn = self.pool.connect()
 
     def tearDown(self) -> None:
-        # Drop the connection before the interpreter shuts down (see the note
-        # in PostgresConnectionTestCase.tearDown).
+        # Drop the connection (returning it to the pool) and the pool before the
+        # interpreter shuts down (see the note in
+        # PostgresConnectionTestCase.tearDown).
         del self.conn
+        del self.pool
 
     # -- small helpers ------------------------------------------------------
 
@@ -1046,10 +1057,12 @@ class PostgresErrorMappingTestCase(unittest.TestCase):
     """
 
     def setUp(self) -> None:
-        self.conn = postgres.connect(_build_dsn())
+        self.pool = postgres.ConnectionPool(_build_dsn())
+        self.conn = self.pool.connect()
 
     def tearDown(self) -> None:
         del self.conn
+        del self.pool
 
     def _exec_commit(self, sql: str) -> None:
         """Run a single statement and commit it (its own transaction)."""
@@ -1102,12 +1115,14 @@ class PostgresErrorMappingTestCase(unittest.TestCase):
     "requires Postgres reachable on libpq's default host",
 )
 class PostgresDefaultHostTestCase(unittest.TestCase):
-    """Covers the libpq default-host fixup in ``connect``.
+    """Covers the libpq default-host fixup.
 
     When the DSN omits ``host=``, ``tokio-postgres`` would default to localhost,
     but Synapse wants libpq's default (honouring ``PGHOST`` / the compiled-in
-    socket dir). This only runs when the test Postgres is actually reachable on
-    that default host, so it's guarded separately from the main suite.
+    socket dir). The fixup lives in the pool's connection manager, so checking a
+    connection out of a pool built from a host-less DSN exercises it. This only
+    runs when the test Postgres is actually reachable on that default host, so
+    it's guarded separately from the main suite.
     """
 
     def test_connect_without_host_uses_libpq_default(self) -> None:
@@ -1120,13 +1135,15 @@ class PostgresDefaultHostTestCase(unittest.TestCase):
             parts.append(f"port={POSTGRES_PORT}")
         if POSTGRES_PASSWORD is not None:
             parts.append(f"password={POSTGRES_PASSWORD}")
-        conn = postgres.connect(" ".join(parts))
+        pool = postgres.ConnectionPool(" ".join(parts))
+        conn = pool.connect()
         try:
             self.assertEqual(
                 run_interaction(conn, lambda cursor: _select_one(cursor)), (1,)
             )
         finally:
             del conn
+            del pool
 
 
 def _select_one(cursor: Any) -> Optional[list[Any]]: