Add RustPostgresEngine for the native Rust Postgres backend

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) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01W3G4M92AmwSSZCbmtMJU3d

synapse/storage/engines/_base.py

@@ -138,9 +138,15 @@ class BaseDatabaseEngine(Generic[ConnectionType, CursorType], metaclass=abc.ABCM
 
         This is not provided by DBAPI2, and so needs engine-specific support.
 
-        Any ongoing transaction is committed before executing the script in its own
-        transaction. The script transaction is left open and it is the responsibility of
-        the caller to commit it.
+        Transaction handling differs by engine. On the psycopg2 and sqlite3
+        engines any ongoing transaction is committed before the script runs in
+        its own transaction, so in a sequence of scripts (as `prepare_database`
+        applies schema/delta files) each one's work is durable once the next
+        starts. The Rust engine instead runs the script in the connection's
+        ongoing transaction, so a whole sequence is applied either completely
+        or not at all (a mid-sequence failure rolls everything back to a
+        consistent, re-runnable state). In all cases the transaction is left
+        open and it is the responsibility of the caller to commit it.
         """
         ...
 

synapse/storage/engines/postgres_rust.py

@@ -0,0 +1,122 @@
+# 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 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.<Error>`.
+        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 — intentionally *without* psycopg2's
+        # `COMMIT; BEGIN TRANSACTION;` prefix, so a whole schema-upgrade run
+        # commits atomically at the end (see BaseDatabaseEngine.executescript).
+        script = script.replace(
+            AUTO_INCREMENT_PRIMARY_KEYPLACEHOLDER,
+            "BIGINT PRIMARY KEY GENERATED ALWAYS AS IDENTITY",
+        )
+        cursor.executescript(script)

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:
+# <https://www.gnu.org/licenses/agpl-3.0.html>.
+
+"""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.<Error>`, 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()