diff --git a/synapse/storage/engines/postgres_rust.py b/synapse/storage/engines/postgres_rust.py index e95dc437eb..43024d89b0 100644 --- a/synapse/storage/engines/postgres_rust.py +++ b/synapse/storage/engines/postgres_rust.py @@ -31,10 +31,6 @@ 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 @@ -43,6 +39,7 @@ from typing import TYPE_CHECKING, Any, Mapping from synapse.storage.engines._base import ( AUTO_INCREMENT_PRIMARY_KEYPLACEHOLDER, IncorrectDatabaseSetup, + IsolationLevel, ) from synapse.storage.engines.postgres_base import PostgresEngine from synapse.storage.types import Connection, Cursor @@ -60,6 +57,18 @@ _RETRYABLE_PGCODES = ("40001", "40P01") class RustPostgresEngine(PostgresEngine[Connection, Cursor]): """A :class:`PostgresEngine` that talks to the Rust backend's shim.""" + # SQL isolation-level names for each `IsolationLevel`. The shim has no + # psycopg2-style `set_isolation_level`, so a per-transaction override is + # applied as a `SET SESSION CHARACTERISTICS` statement (see + # `attempt_to_set_isolation_level`). REPEATABLE READ is the default that the + # pool applies at connection setup, matching `PostgresEngine`. + _ISOLATION_LEVEL_SQL: Mapping[int, str] = { + IsolationLevel.READ_COMMITTED: "READ COMMITTED", + IsolationLevel.REPEATABLE_READ: "REPEATABLE READ", + IsolationLevel.SERIALIZABLE: "SERIALIZABLE", + } + _DEFAULT_ISOLATION_LEVEL_SQL = "REPEATABLE READ" + def __init__(self, database_config: Mapping[str, Any]): # The module is the Rust backend's DBAPI2 exception hierarchy # (OperationalError, DatabaseError, IntegrityError, …); the transaction @@ -108,12 +117,29 @@ class RustPostgresEngine(PostgresEngine[Connection, Cursor]): 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" - ) + # psycopg2 has `conn.set_isolation_level`; the shim does not, so set the + # session's default transaction isolation directly. It applies to the + # next transaction the caller opens — `new_transaction` sets it before + # the transaction and resets it (`isolation_level=None`) afterwards, so + # the override is scoped to that one transaction. `isolation_level` is a + # fixed `IsolationLevel`, so the interpolated name is not user-controlled. + if isolation_level is None: + level_sql = self._DEFAULT_ISOLATION_LEVEL_SQL + else: + level_sql = self._ISOLATION_LEVEL_SQL[isolation_level] + + # `SET SESSION CHARACTERISTICS` only affects *subsequent* transactions, + # so commit it (SET is transactional — a rollback would undo it) before + # the caller's transaction begins. + cursor = conn.cursor() + try: + cursor.execute( + "SET SESSION CHARACTERISTICS AS TRANSACTION ISOLATION LEVEL " + + level_sql + ) + finally: + cursor.close() + conn.commit() @staticmethod def executescript(cursor: Any, script: str) -> None: diff --git a/tests/storage/test_rust_engine.py b/tests/storage/test_rust_engine.py index b624362120..0665bfdac9 100644 --- a/tests/storage/test_rust_engine.py +++ b/tests/storage/test_rust_engine.py @@ -15,6 +15,7 @@ from typing import Any from synapse.storage.engines import PostgresEngine, Psycopg2Engine, create_engine +from synapse.storage.engines._base import IsolationLevel from synapse.storage.engines.postgres_rust import RustPostgresEngine from synapse.synapse_rust.database import postgres @@ -81,9 +82,40 @@ class RustPostgresEngineTestCase(unittest.TestCase): # 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) + def test_isolation_level_sets_session_characteristics(self) -> None: + # The shim has no psycopg2 `set_isolation_level`, so the engine issues a + # `SET SESSION CHARACTERISTICS` statement (and commits it) instead. Check + # the level mapping, and that `None` resets to the default (REPEATABLE + # READ) — see `attempt_to_set_isolation_level`. + executed: list[str] = [] + + class FakeCursor: + def execute(self, sql: str) -> None: + executed.append(sql) + + def close(self) -> None: + pass + + class FakeConn: + def cursor(self) -> "FakeCursor": + return FakeCursor() + + def commit(self) -> None: + executed.append("COMMIT") + + conn = FakeConn() + self.engine.attempt_to_set_isolation_level(conn, IsolationLevel.SERIALIZABLE) + self.engine.attempt_to_set_isolation_level(conn, None) + + self.assertEqual( + executed, + [ + "SET SESSION CHARACTERISTICS AS TRANSACTION ISOLATION LEVEL SERIALIZABLE", + "COMMIT", + "SET SESSION CHARACTERISTICS AS TRANSACTION ISOLATION LEVEL REPEATABLE READ", + "COMMIT", + ], + ) def test_create_engine_selects_rust_only_when_opted_in(self) -> None: # Default: the psycopg2 engine.