From da873902abea954e7eeab469eeb63a401e7a7741 Mon Sep 17 00:00:00 2001 From: Erik Johnston Date: Thu, 2 Jul 2026 15:31:58 +0000 Subject: [PATCH] Add RustPostgresEngine for the native Rust Postgres backend MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A database engine that drives the Rust Connection/Cursor shim instead of psycopg2. It subclasses PostgresEngine to reuse the pure SQL-generation and configuration behaviour, and overrides only the parts that touch a live connection or are wired to psycopg2 internals: - `module` points at the Rust backend's DBAPI2 exception hierarchy, which the transaction driver catches on; - `convert_param_style` rewrites `?` to `$1, $2, ...` (the shim binds libpq positional placeholders, not psycopg2's `%s`); - `in_transaction` / `is_connection_closed` / `attempt_to_set_autocommit` call the shim's own methods; - `is_deadlock` matches the Rust `DatabaseError` and its `pgcode`; - `executescript` uses the shim's multi-statement primitive; - `on_new_connection` is a no-op — session setup belongs in the Rust pool. Per-transaction isolation-level overrides raise NotImplementedError for now, and `check_database` / `server_version` still read psycopg2 attributes; the engine is not yet selectable via `create_engine`, so neither is reached. Wiring it into `make_pool` (which also feeds session config to the pool) is the next step. Tested directly: param-style rewriting, deadlock/pgcode matching, the module pointer, and — against a live shim connection — in_transaction, is_closed and autocommit. Co-Authored-By: Claude Opus 4.8 (1M context) Claude-Session: https://claude.ai/code/session_01W3G4M92AmwSSZCbmtMJU3d --- synapse/storage/engines/postgres_rust.py | 120 +++++++++++++++++++++ tests/storage/test_rust_engine.py | 132 +++++++++++++++++++++++ 2 files changed, 252 insertions(+) create mode 100644 synapse/storage/engines/postgres_rust.py create mode 100644 tests/storage/test_rust_engine.py diff --git a/synapse/storage/engines/postgres_rust.py b/synapse/storage/engines/postgres_rust.py new file mode 100644 index 0000000000..09310d4db0 --- /dev/null +++ b/synapse/storage/engines/postgres_rust.py @@ -0,0 +1,120 @@ +# 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 database engine for the native Rust ``tokio-postgres`` backend. + +This drives the Rust ``Connection`` / ``Cursor`` shim +(:mod:`synapse.synapse_rust.database.postgres`) rather than psycopg2. It reuses +:class:`PostgresEngine` for everything that is pure SQL generation or +configuration (``supports_using_any_list``, ``lock_table``, ``row_id_name``, …) +and overrides only the parts that touch a live connection or that are wired to +psycopg2 internals: + + - the DBAPI2 exception ``module`` — the Rust backend has its own hierarchy; + - ``convert_param_style`` — the shim binds ``$1, $2, …`` placeholders, not + psycopg2's ``%s``; + - ``in_transaction`` / ``is_connection_closed`` / ``attempt_to_set_autocommit`` + — served by the shim's own methods; + - ``is_deadlock`` — matches the Rust ``DatabaseError`` and its ``pgcode``; + - ``executescript`` — uses the shim's dedicated multi-statement primitive. + +Per-connection session setup (isolation level, ``synchronous_commit``, +``statement_timeout``) lives in the Rust connection pool rather than in +``on_new_connection``, so that hook is a no-op here. + +Not yet adapted (the engine is not yet selectable via ``create_engine``, so +these are not reached): ``check_database`` / ``server_version`` still read +psycopg2 connection attributes, and per-transaction isolation-level overrides +are unimplemented. Both are follow-ups for the full ``make_pool`` wiring. +""" + +import logging +from typing import TYPE_CHECKING, Any, Mapping + +from synapse.storage.engines._base import AUTO_INCREMENT_PRIMARY_KEYPLACEHOLDER +from synapse.storage.engines.postgres import PostgresEngine +from synapse.synapse_rust.database import postgres + +if TYPE_CHECKING: + from synapse.storage.database import LoggingDatabaseConnection + +logger = logging.getLogger(__name__) + +# Deadlock / serialization-failure SQLSTATEs that Synapse retries. +_RETRYABLE_PGCODES = ("40001", "40P01") + + +class RustPostgresEngine(PostgresEngine): + """A :class:`PostgresEngine` that talks to the Rust backend's shim.""" + + def __init__(self, database_config: Mapping[str, Any]): + super().__init__(database_config) + # Route the DBAPI2 exception hierarchy (OperationalError, DatabaseError, + # IntegrityError, …) to the Rust backend's classes; the transaction + # driver catches `engine.module.`. + self.module = postgres + + def convert_param_style(self, sql: str) -> str: + # The shim binds positional `$1, $2, ...` placeholders (like libpq), + # not psycopg2's `%s`. Rewrite `?` left-to-right, matching the Rust-side + # `convert_placeholders`; callers must parameterise rather than embed a + # literal `?`. + out = [] + n = 0 + for ch in sql: + if ch == "?": + n += 1 + out.append(f"${n}") + else: + out.append(ch) + return "".join(out) + + def on_new_connection(self, db_conn: "LoggingDatabaseConnection") -> None: + # No-op: per-connection session setup happens in the Rust connection + # pool's connection manager, not here. + pass + + def is_deadlock(self, error: Exception) -> bool: + if isinstance(error, postgres.DatabaseError): + return getattr(error, "pgcode", None) in _RETRYABLE_PGCODES + return False + + def is_connection_closed(self, conn: Any) -> bool: + return bool(conn.is_closed()) + + def in_transaction(self, conn: Any) -> bool: + return bool(conn.in_transaction()) + + def attempt_to_set_autocommit(self, conn: Any, autocommit: bool) -> None: + conn.set_autocommit(autocommit) + + def attempt_to_set_isolation_level( + self, conn: Any, isolation_level: int | None + ) -> None: + # Per-transaction isolation overrides are not implemented for the shim + # yet; the connection's default level is set when the pool opens it. + raise NotImplementedError( + "per-transaction isolation levels are not supported by the Rust " + "Postgres backend yet" + ) + + @staticmethod + def executescript(cursor: Any, script: str) -> None: + # Use the shim's dedicated multi-statement primitive rather than + # psycopg2's "just execute it" behaviour. The script runs in the + # connection's current transaction (opened lazily) and is left open for + # the caller to commit. + script = script.replace( + AUTO_INCREMENT_PRIMARY_KEYPLACEHOLDER, + "BIGINT PRIMARY KEY GENERATED ALWAYS AS IDENTITY", + ) + cursor.executescript(script) diff --git a/tests/storage/test_rust_engine.py b/tests/storage/test_rust_engine.py new file mode 100644 index 0000000000..98cf70f124 --- /dev/null +++ b/tests/storage/test_rust_engine.py @@ -0,0 +1,132 @@ +# 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 :class:`synapse.storage.engines.postgres_rust.RustPostgresEngine`.""" + +from typing import Any + +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) + + +class RustPostgresEngineTestCase(unittest.TestCase): + """The engine's connection-independent behaviour (no database needed).""" + + def setUp(self) -> None: + self.engine = RustPostgresEngine({}) + + def test_module_is_the_rust_backend(self) -> None: + # The transaction driver catches `engine.module.`, so the module + # must be the Rust backend's DBAPI2 hierarchy, not psycopg2's. + self.assertIs(self.engine.module, postgres) + + def test_convert_param_style_rewrites_to_dollar_placeholders(self) -> None: + self.assertEqual( + self.engine.convert_param_style("SELECT * FROM t WHERE a = ? AND b = ?"), + "SELECT * FROM t WHERE a = $1 AND b = $2", + ) + # No placeholders: unchanged. Past nine: decimal, not a single digit. + self.assertEqual(self.engine.convert_param_style("SELECT 1"), "SELECT 1") + self.assertTrue( + self.engine.convert_param_style("VALUES " + "(?)" * 11).endswith("($11)") + ) + + def test_is_deadlock_matches_rust_database_error_pgcodes(self) -> None: + for pgcode in ("40001", "40P01"): + err = postgres.DatabaseError("boom") + err.pgcode = pgcode + self.assertTrue(self.engine.is_deadlock(err)) + + # A DatabaseError with some other pgcode is not a deadlock. + other = postgres.DatabaseError("nope") + other.pgcode = "23505" + self.assertFalse(self.engine.is_deadlock(other)) + + # Neither is an unrelated exception. + self.assertFalse(self.engine.is_deadlock(ValueError("unrelated"))) + + def test_isolation_level_override_not_yet_supported(self) -> None: + with self.assertRaises(NotImplementedError): + self.engine.attempt_to_set_isolation_level(object(), None) + + +@skip_unless( + bool(USE_POSTGRES_FOR_TESTS), "requires a Postgres server (set SYNAPSE_POSTGRES)" +) +class RustPostgresEngineConnectionTestCase(unittest.TestCase): + """The engine's connection-touching checks, against a live shim connection.""" + + def setUp(self) -> None: + self.engine = RustPostgresEngine({}) + self.pool = postgres.ConnectionPool(_build_dsn()) + self.conn = self.pool.connect() + + def tearDown(self) -> None: + del self.conn + self.pool.close() + + def _exec(self, sql: str) -> Any: + cursor = self.conn.cursor() + cursor.execute(sql) + return cursor + + def test_in_transaction_tracks_the_transaction(self) -> None: + self.assertFalse(self.engine.in_transaction(self.conn)) + # The first statement lazily opens a transaction. + self._exec("SELECT 1") + self.assertTrue(self.engine.in_transaction(self.conn)) + self.conn.commit() + self.assertFalse(self.engine.in_transaction(self.conn)) + + def test_is_connection_closed(self) -> None: + self.assertFalse(self.engine.is_connection_closed(self.conn)) + self.conn.close() + self.assertTrue(self.engine.is_connection_closed(self.conn)) + + def test_attempt_to_set_autocommit(self) -> None: + # In autocommit mode the shim issues no implicit BEGIN, so a statement + # does not open a transaction. + self.engine.attempt_to_set_autocommit(self.conn, True) + self._exec("SELECT 1") + self.assertFalse(self.engine.in_transaction(self.conn)) + + # Turning it back off restores the transactional default. + self.engine.attempt_to_set_autocommit(self.conn, False) + self._exec("SELECT 1") + self.assertTrue(self.engine.in_transaction(self.conn)) + self.conn.rollback()