mirror of
https://github.com/element-hq/synapse.git
synced 2026-08-15 04:50:25 +00:00
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) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01W3G4M92AmwSSZCbmtMJU3d
This commit is contained in:
Erik Johnston
co-authored by
Claude Opus 4.8
Claude Opus 4.8
parent
ab33ca232f
commit
969a473d58
@@ -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:
|
||||
# <https://www.gnu.org/licenses/agpl-3.0.html>.
|
||||
|
||||
"""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)
|
||||
@@ -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:
|
||||
# <https://www.gnu.org/licenses/agpl-3.0.html>.
|
||||
|
||||
"""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()
|
||||
Reference in New Issue
Block a user