synapse/tests/storage/test_rust_engine.py

# 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 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

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_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.
        self.assertIsInstance(create_engine({"name": "psycopg2"}), Psycopg2Engine)

        # With the opt-in flag: the Rust engine — still a PostgresEngine, so the
        # storage layer's `isinstance(engine, PostgresEngine)` dialect checks hold.
        engine = create_engine({"name": "psycopg2", "use_rust_driver": True})
        self.assertIsInstance(engine, RustPostgresEngine)
        self.assertIsInstance(engine, PostgresEngine)
        self.assertNotIsInstance(engine, Psycopg2Engine)


@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_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.
        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()