From 267238d72d7cf8a7a9fd2b32d0643cc6cdced111 Mon Sep 17 00:00:00 2001 From: Erik Johnston Date: Fri, 3 Jul 2026 10:43:48 +0000 Subject: [PATCH] Support per-transaction isolation levels on the Rust backend MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `new_transaction` sets a transaction's isolation level (for the receipts, push-actions and purge paths, which ask for a specific level) by calling `engine.attempt_to_set_isolation_level` before the transaction and resetting it after. The Rust engine raised `NotImplementedError` there, so those transactions errored out on the Rust backend. The shim has no psycopg2-style `conn.set_isolation_level`, so implement it in SQL: run `SET SESSION CHARACTERISTICS AS TRANSACTION ISOLATION LEVEL ` on the connection. That sets the session's default for *subsequent* transactions, so it's committed before the caller's transaction begins (SET is transactional); `new_transaction` resets it to the default (REPEATABLE READ, which the pool already applies at connection setup) afterwards, scoping the override to the one transaction — matching psycopg2's behaviour. The level name comes from a fixed `IsolationLevel` map, not caller input. Fixes the isolation-level failures in tests.storage.test_event_push_actions, tests.storage.test_purge and tests.storage.test_receipts on the Rust backend (20/20 pass); psycopg2 and sqlite are unaffected. Co-Authored-By: Claude Opus 4.8 (1M context) Claude-Session: https://claude.ai/code/session_01W3G4M92AmwSSZCbmtMJU3d --- synapse/storage/database.py | 4 ++ synapse/storage/engines/postgres_rust.py | 57 ++++++++++++++---- tests/storage/test_rust_engine.py | 73 +++++++++++++++++++++++- 3 files changed, 121 insertions(+), 13 deletions(-) diff --git a/synapse/storage/database.py b/synapse/storage/database.py index 1dee889067..59d7eba21c 100644 --- a/synapse/storage/database.py +++ b/synapse/storage/database.py @@ -1131,6 +1131,10 @@ class DatabasePool: i.e. outside of a transaction. This is useful for transaction that are only a single query. Currently only affects postgres. isolation_level: Set the server isolation level for this transaction. + Note: resetting the level afterwards rolls back anything `func` + left uncommitted (psycopg2's `set_isolation_level` semantics), + so `func` must commit its own work — the connection pool's + commit-on-success does not apply to it. kwargs: named args to pass to `func` Returns: diff --git a/synapse/storage/engines/postgres_rust.py b/synapse/storage/engines/postgres_rust.py index 01c0886372..dc67651732 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,40 @@ 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] + + # psycopg2's `set_isolation_level` implicitly rolls back any pending + # transaction before changing the session; mirror that. This matters + # when the reset runs from `runWithConnection`'s `finally` after the + # transaction function raised without `new_transaction` rolling back + # (anything but an OperationalError/deadlock): without the rollback the + # `SET` would either fail with 25P02 on an aborted transaction (masking + # the original error) or — worse — the commit below would persist the + # half-finished writes. + conn.rollback() + + # `SET SESSION CHARACTERISTICS` only affects *subsequent* transactions, + # so commit it (SET is transactional — a rollback would undo it) before + # the caller's transaction begins. The rollback above ensures the + # transaction being committed contains nothing but the `SET`. + 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..951b78607a 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,46 @@ 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 — + # rolling back any pending transaction first, as psycopg2 does. 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") + + def rollback(self) -> None: + executed.append("ROLLBACK") + + conn = FakeConn() + self.engine.attempt_to_set_isolation_level(conn, IsolationLevel.SERIALIZABLE) + self.engine.attempt_to_set_isolation_level(conn, None) + + self.assertEqual( + executed, + [ + "ROLLBACK", + "SET SESSION CHARACTERISTICS AS TRANSACTION ISOLATION LEVEL SERIALIZABLE", + "COMMIT", + "ROLLBACK", + "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. @@ -130,6 +168,35 @@ class RustPostgresEngineConnectionTestCase(unittest.TestCase): self.conn.close() self.assertTrue(self.engine.is_connection_closed(self.conn)) + def test_isolation_level_reset_rolls_back_pending_writes(self) -> None: + # `attempt_to_set_isolation_level` runs from `runWithConnection`'s + # `finally` after a transaction function raised, when the transaction + # may still be open with uncommitted writes: it must roll those back + # (as psycopg2's `set_isolation_level` does), not commit them. + self._exec("CREATE TEMPORARY TABLE isolation_reset_test (x INT)") + self.conn.commit() + + self._exec("INSERT INTO isolation_reset_test VALUES (1)") + self.engine.attempt_to_set_isolation_level(self.conn, None) + + cursor = self._exec("SELECT COUNT(*) FROM isolation_reset_test") + self.assertEqual(cursor.fetch_one()[0], 0) + self.conn.rollback() + + def test_isolation_level_reset_survives_aborted_transaction(self) -> None: + # Resetting the isolation level on a server-side-aborted transaction + # must not raise 25P02 ("current transaction is aborted"), which would + # mask the error that aborted it. + with self.assertRaises(postgres.DatabaseError): + self._exec("SELECT no_such_column FROM nonexistent_table") + + self.engine.attempt_to_set_isolation_level(self.conn, None) + + # The connection is usable again afterwards. + cursor = self._exec("SELECT 1") + self.assertEqual(cursor.fetch_one()[0], 1) + self.conn.rollback() + 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.