From 969a473d5883cf189f307dda6c041bbd2728c366 Mon Sep 17 00:00:00 2001 From: Erik Johnston Date: Thu, 2 Jul 2026 16:01:24 +0000 Subject: [PATCH] Add a DBAPI2 adapter over the Rust Postgres shim Synapse's LoggingTransaction drives a cursor through the DBAPI2 spelling (fetchone/fetchmany/fetchall, iteration, rowcount/description as properties), but the Rust shim cursor exposes fetch_one/fetch_all/fetch_next_batch and rowcount()/description() as methods. Rather than reshape the Rust API, add thin Python `Connection`/`Cursor` wrappers that present the DBAPI2 shape and delegate to the shim: - `Cursor` maps fetchone/fetchmany/fetchall/__iter__ and rowcount/description onto the shim, and tracks exhaustion so fetching past the end keeps returning "no more rows" (the shim raises instead); - `Connection.cursor()` returns the adapter cursor, and the transaction-control and engine-facing methods (commit/rollback/close, set_autocommit, is_closed, in_transaction) delegate straight through. An end-to-end test drives a real LoggingTransaction backed by the adapter and RustPostgresEngine: `?` placeholders are converted to `$n`, a query runs, and rows come back via fetchone/iteration/description. Not handled yet: execute_batch / execute_values (psycopg2 extras that LoggingTransaction calls directly for PostgresEngine) still need a routing change to reach a shim-backed implementation. Co-Authored-By: Claude Opus 4.8 (1M context) Claude-Session: https://claude.ai/code/session_01W3G4M92AmwSSZCbmtMJU3d --- synapse/storage/rust_dbapi.py | 156 +++++++++++++++++++++++++++++++ tests/storage/test_rust_dbapi.py | 131 ++++++++++++++++++++++++++ 2 files changed, 287 insertions(+) create mode 100644 synapse/storage/rust_dbapi.py create mode 100644 tests/storage/test_rust_dbapi.py diff --git a/synapse/storage/rust_dbapi.py b/synapse/storage/rust_dbapi.py new file mode 100644 index 0000000000..3aa696d722 --- /dev/null +++ b/synapse/storage/rust_dbapi.py @@ -0,0 +1,156 @@ +# This file is licensed under the Affero General Public License (AGPL) version 3. +# +# Copyright (C) 2026 Element Creations Ltd +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Affero General Public License as +# published by the Free Software Foundation, either version 3 of the +# License, or (at your option) any later version. +# +# See the GNU Affero General Public License for more details: +# . + +"""A thin DBAPI2 adapter over the native Rust Postgres shim. + +The Rust ``Connection`` / ``Cursor`` shim +(:mod:`synapse.synapse_rust.database.postgres`) is close to DBAPI2 but not +identical: its cursor exposes ``fetch_one`` / ``fetch_all`` / +``fetch_next_batch`` and ``rowcount()`` / ``description()`` as *methods*, whereas +Synapse's :class:`~synapse.storage.database.LoggingTransaction` drives a cursor +through the DBAPI2 spelling — ``fetchone`` / ``fetchmany`` / ``fetchall``, +iteration, and ``rowcount`` / ``description`` as *properties*. + +Rather than reshape the Rust API, these small wrappers present the DBAPI2 shape +Synapse expects and delegate to the shim. :class:`Connection` also keeps the +connection-level methods the database engine calls (``in_transaction``, +``is_closed``, ``set_autocommit``) so a wrapped connection is a drop-in for the +raw one. + +Not handled here: ``execute_batch`` / ``execute_values`` (psycopg2 extras that +``LoggingTransaction`` invokes directly for ``PostgresEngine``) still need a +routing change in ``LoggingTransaction`` to reach a shim-backed implementation; +that is a separate follow-up. +""" + +from typing import TYPE_CHECKING, Any, Iterator, Sequence + +if TYPE_CHECKING: + from synapse.storage.types import SQLQueryParameters + + +class Cursor: + """A DBAPI2 cursor wrapping a Rust shim cursor.""" + + # DBAPI2 default number of rows `fetchmany` returns when no size is given. + arraysize = 1 + + def __init__(self, cursor: Any) -> None: + self._cursor = cursor + # The shim raises if a result set is fetched past exhaustion, whereas + # DBAPI2 wants further fetches to keep returning "no more rows". Track + # exhaustion here so repeated / mixed fetches stay well-behaved. + self._exhausted = False + + def execute(self, sql: str, parameters: "SQLQueryParameters" = ()) -> None: + # DBAPI2 passes an empty sequence when there are no parameters; the shim + # wants `None` in that case (and a list of values otherwise). + self._exhausted = False + self._cursor.execute(sql, list(parameters) if parameters else None) + + def executemany(self, sql: str, seq_of_parameters: Sequence[Any]) -> None: + self._exhausted = False + self._cursor.executemany(sql, [list(p) for p in seq_of_parameters]) + + def _next(self) -> Any: + """Fetch one row, or `None` once the result set is exhausted.""" + if self._exhausted: + return None + row = self._cursor.fetch_one() + if row is None: + self._exhausted = True + return row + + def fetchone(self) -> Any: + return self._next() + + def fetchmany(self, size: int | None = None) -> list[Any]: + # DBAPI2: return at most `size` rows (defaulting to `arraysize`). The + # shim's `fetch_next_batch` is a size *hint* rather than a hard limit, so + # pull rows one at a time to honour the exact contract. + if size is None: + size = self.arraysize + rows = [] + for _ in range(size): + row = self._next() + if row is None: + break + rows.append(row) + return rows + + def fetchall(self) -> list[Any]: + if self._exhausted: + return [] + rows = self._cursor.fetch_all() + self._exhausted = True + return rows + + def __iter__(self) -> Iterator[Any]: + while True: + row = self._next() + if row is None: + return + yield row + + @property + def rowcount(self) -> int: + return self._cursor.rowcount() + + @property + def description(self) -> Any: + return self._cursor.description() + + def close(self) -> None: + self._cursor.close() + + +class Connection: + """A DBAPI2 connection wrapping a Rust shim connection. + + ``cursor()`` returns a DBAPI2 :class:`Cursor`; the transaction-control and + engine-facing methods delegate straight to the shim. + """ + + def __init__(self, conn: Any) -> None: + self._conn = conn + + def cursor(self) -> Cursor: + return Cursor(self._conn.cursor()) + + def commit(self) -> None: + self._conn.commit() + + def rollback(self) -> None: + self._conn.rollback() + + def close(self) -> None: + self._conn.close() + + # -- engine-facing methods (see RustPostgresEngine) --------------------- + + def set_autocommit(self, autocommit: bool) -> None: + self._conn.set_autocommit(autocommit) + + def is_closed(self) -> bool: + return bool(self._conn.is_closed()) + + def in_transaction(self) -> bool: + return bool(self._conn.in_transaction()) + + # -- context-manager protocol (matches the shim / psycopg2) ------------- + + def __enter__(self) -> "Connection": + self._conn.__enter__() + return self + + def __exit__(self, exc_type: Any, exc: Any, tb: Any) -> Any: + return self._conn.__exit__(exc_type, exc, tb) diff --git a/tests/storage/test_rust_dbapi.py b/tests/storage/test_rust_dbapi.py new file mode 100644 index 0000000000..812e81022f --- /dev/null +++ b/tests/storage/test_rust_dbapi.py @@ -0,0 +1,131 @@ +# This file is licensed under the Affero General Public License (AGPL) version 3. +# +# Copyright (C) 2026 Element Creations Ltd +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Affero General Public License as +# published by the Free Software Foundation, either version 3 of the +# License, or (at your option) any later version. +# +# See the GNU Affero General Public License for more details: +# . + +"""Tests for the DBAPI2 adapter over the Rust Postgres shim +(:mod:`synapse.storage.rust_dbapi`), including driving a real +``LoggingTransaction`` through it.""" + +from synapse.storage import rust_dbapi +from synapse.storage.database import LoggingDatabaseConnection +from synapse.storage.engines.postgres_rust import RustPostgresEngine +from synapse.synapse_rust.database import postgres + +from tests import unittest +from tests.unittest import skip_unless +from tests.utils import ( + POSTGRES_BASE_DB, + POSTGRES_HOST, + POSTGRES_PASSWORD, + POSTGRES_PORT, + POSTGRES_USER, + USE_POSTGRES_FOR_TESTS, +) + + +def _build_dsn() -> str: + """Build a libpq keyword/value connection string from the test config.""" + + parts = [f"dbname={POSTGRES_BASE_DB}"] + if POSTGRES_USER is not None: + parts.append(f"user={POSTGRES_USER}") + if POSTGRES_HOST is not None: + parts.append(f"host={POSTGRES_HOST}") + if POSTGRES_PORT is not None: + parts.append(f"port={POSTGRES_PORT}") + if POSTGRES_PASSWORD is not None: + parts.append(f"password={POSTGRES_PASSWORD}") + return " ".join(parts) + + +@skip_unless( + bool(USE_POSTGRES_FOR_TESTS), "requires a Postgres server (set SYNAPSE_POSTGRES)" +) +class RustDBAPIAdapterTestCase(unittest.TestCase): + """The adapter presents the DBAPI2 shape over the shim.""" + + def setUp(self) -> None: + self._pool = postgres.ConnectionPool(_build_dsn()) + self.conn = rust_dbapi.Connection(self._pool.connect()) + + def tearDown(self) -> None: + del self.conn + self._pool.close() + + def test_execute_and_fetchone(self) -> None: + cursor = self.conn.cursor() + # The adapter passes parameters straight through; the shim binds `$n`. + cursor.execute("SELECT $1::int", (7,)) + self.assertEqual(cursor.fetchone(), (7,)) + # Exhausted → None. + self.assertIsNone(cursor.fetchone()) + self.conn.commit() + + def test_fetchall(self) -> None: + cursor = self.conn.cursor() + cursor.execute("SELECT g FROM generate_series(1, 3) AS g ORDER BY g") + self.assertEqual(cursor.fetchall(), [(1,), (2,), (3,)]) + self.conn.commit() + + def test_fetchmany_returns_at_most_size(self) -> None: + cursor = self.conn.cursor() + cursor.execute("SELECT g FROM generate_series(1, 3) AS g ORDER BY g") + self.assertEqual(cursor.fetchmany(2), [(1,), (2,)]) + self.assertEqual(cursor.fetchmany(2), [(3,)]) + self.assertEqual(cursor.fetchmany(2), []) + self.conn.commit() + + def test_iteration(self) -> None: + cursor = self.conn.cursor() + cursor.execute("SELECT g FROM generate_series(1, 3) AS g ORDER BY g") + self.assertEqual(list(cursor), [(1,), (2,), (3,)]) + self.conn.commit() + + def test_description_exposes_column_names(self) -> None: + cursor = self.conn.cursor() + cursor.execute("SELECT 1 AS a, 2 AS b") + assert cursor.description is not None + self.assertEqual([col[0] for col in cursor.description], ["a", "b"]) + self.conn.commit() + + def test_rowcount_and_executemany(self) -> None: + cursor = self.conn.cursor() + cursor.execute("CREATE TEMP TABLE t (id int)") + cursor.executemany("INSERT INTO t VALUES ($1)", [(1,), (2,), (3,)]) + self.assertEqual(cursor.rowcount, 3) + cursor.execute("SELECT id FROM t ORDER BY id") + self.assertEqual(cursor.fetchall(), [(1,), (2,), (3,)]) + self.conn.commit() + + def test_drives_a_logging_transaction(self) -> None: + # The whole point: a real LoggingTransaction (which converts `?` to `$n` + # via the engine, then drives the cursor via the DBAPI2 spelling) runs + # unchanged against the adapter. + engine = RustPostgresEngine({}) + db_conn = LoggingDatabaseConnection( + conn=self.conn, + engine=engine, + default_txn_name="test", + server_name="test", + ) + + txn = db_conn.cursor(txn_name="test") + txn.execute("SELECT ?::int + ?::int", (2, 3)) + self.assertEqual(txn.fetchone(), (5,)) + + txn.execute("SELECT g FROM generate_series(1, 2) AS g ORDER BY g") + self.assertEqual(list(txn), [(1,), (2,)]) + + txn.execute("SELECT 1 AS only") + assert txn.description is not None + self.assertEqual(txn.description[0][0], "only") + + db_conn.commit()