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
This commit is contained in:
Erik Johnston
2026-07-11 09:04:57 +00:00
co-authored by Claude Opus 4.8
parent 2a4ad76bda
commit cb7b622656
4 changed files with 284 additions and 82 deletions
+32 -15
View File
@@ -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]]: