From 56e347225267c7fd7f2bd3b6de747724dc361444 Mon Sep 17 00:00:00 2001 From: Erik Johnston Date: Thu, 2 Jul 2026 17:08:11 +0000 Subject: [PATCH] Wire the startup path for the Rust Postgres backend MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Startup opens a bootstrap connection via `make_conn` (outside the runtime pool) and validates it with `engine.check_database` before `prepare_database` runs. Neither worked for the Rust backend: the Rust module has no `module.connect`, and check_database/server_version were NotImplementedError stubs. Fill them in: - `rust_dbapi.build_dsn` turns the libpq-style `args` into a DSN string, and `rust_dbapi.connect` opens a standalone connection (a pool of one, kept alive by the returned Connection) for one-off bootstrap use. - `make_conn` routes the Rust engine through `rust_dbapi.connect` (with the engine's synchronous_commit / statement_timeout) instead of `engine.module.connect`; psycopg2 and sqlite are unchanged. - `RustPostgresEngine.check_database` reads the server version over a cursor (`SHOW server_version_num`) rather than psycopg2's `conn.server_version`, and applies the same version / encoding / collation / ctype checks; `server_version` is derived from the cached value. Tested: build_dsn quoting; and, against a live Postgres, make_conn + check_database + server_version + a query. psycopg2 and sqlite homeserver boots (make_conn → check_database) remain green. Co-Authored-By: Claude Opus 4.8 (1M context) Claude-Session: https://claude.ai/code/session_01W3G4M92AmwSSZCbmtMJU3d --- synapse/storage/database.py | 18 ++- synapse/storage/engines/postgres_rust.py | 67 ++++++++-- synapse/storage/rust_dbapi.py | 158 ++++++++++++++++++++++- tests/storage/test_rust_dbapi.py | 126 +++++++++++++++++- 4 files changed, 354 insertions(+), 15 deletions(-) diff --git a/synapse/storage/database.py b/synapse/storage/database.py index 023014276b..3c6e9f81be 100644 --- a/synapse/storage/database.py +++ b/synapse/storage/database.py @@ -192,7 +192,23 @@ def make_conn( for k, v in db_config.config.get("args", {}).items() if not k.startswith("cp_") } - native_db_conn = engine.module.connect(**db_params) + + from synapse.storage.engines.postgres_rust import RustPostgresEngine + + native_db_conn: Any + if isinstance(engine, RustPostgresEngine): + # The Rust backend has no `module.connect`; open a standalone (pool-of-one) + # connection from the same libpq args, with the engine's session settings. + from synapse.storage import rust_dbapi + + native_db_conn = rust_dbapi.connect( + rust_dbapi.build_dsn(db_params), + synchronous_commit=engine.synchronous_commit, + statement_timeout_ms=engine.statement_timeout, + ) + else: + native_db_conn = engine.module.connect(**db_params) + db_conn = LoggingDatabaseConnection( conn=native_db_conn, engine=engine, diff --git a/synapse/storage/engines/postgres_rust.py b/synapse/storage/engines/postgres_rust.py index 7d5b100714..01c0886372 100644 --- a/synapse/storage/engines/postgres_rust.py +++ b/synapse/storage/engines/postgres_rust.py @@ -40,7 +40,10 @@ 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._base import ( + AUTO_INCREMENT_PRIMARY_KEYPLACEHOLDER, + IncorrectDatabaseSetup, +) from synapse.storage.engines.postgres_base import PostgresEngine from synapse.storage.types import Connection, Cursor from synapse.synapse_rust.database import postgres @@ -66,6 +69,7 @@ class RustPostgresEngine(PostgresEngine[Connection, Cursor]): # via `rust_dbapi.connect`), so it doesn't structurally satisfy the # protocol — hence the ignore. super().__init__(postgres, database_config) # type: ignore[arg-type] + self._version: int | None = None # set by check_database def convert_param_style(self, sql: str) -> str: # The shim binds positional `$1, $2, ...` placeholders (like libpq), @@ -128,16 +132,57 @@ class RustPostgresEngine(PostgresEngine[Connection, Cursor]): def check_database( self, db_conn: Any, allow_outdated_version: bool = False ) -> None: - # Startup database validation reads psycopg2 connection attributes - # (server_version, ...) that the shim doesn't expose; adapting it is - # part of wiring the Rust backend into startup (a follow-up). - raise NotImplementedError( - "check_database is not yet implemented for the Rust Postgres backend" - ) + # The shim has no psycopg2-style `conn.server_version`, so read the + # version (and encoding) over a cursor instead. + allow_unsafe_locale = self.config.get("allow_unsafe_locale", False) + + with db_conn.cursor() as cur: + cur.execute("SHOW server_version_num") + self._version = int(cur.fetchone()[0]) + + # Are we on a supported PostgreSQL version? + if not allow_outdated_version and self._version < 140000: + raise RuntimeError("Synapse requires PostgreSQL 14 or above.") + + cur.execute("SHOW SERVER_ENCODING") + rows = cur.fetchall() + if rows and rows[0][0] != "UTF8": + raise IncorrectDatabaseSetup( + "Database has incorrect encoding: '%s' instead of 'UTF8'\n" + "See docs/postgres.md for more information." % (rows[0][0],) + ) + + collation, ctype = self.get_db_locale(cur) + if collation != "C": + logger.warning( + "Database has incorrect collation of %r. Should be 'C'", + collation, + ) + if not allow_unsafe_locale: + raise IncorrectDatabaseSetup( + "Database has incorrect collation of %r. Should be 'C'\n" + "See docs/postgres.md for more information. You can override this check by" + "setting 'allow_unsafe_locale' to true in the database config.", + collation, + ) + + if ctype != "C" and not allow_unsafe_locale: + logger.warning( + "Database has incorrect ctype of %r. Should be 'C'", + ctype, + ) + raise IncorrectDatabaseSetup( + "Database has incorrect ctype of %r. Should be 'C'\n" + "See docs/postgres.md for more information. You can override this check by" + "setting 'allow_unsafe_locale' to true in the database config.", + ctype, + ) @property def server_version(self) -> str: - # As above: depends on the psycopg2 startup path that isn't wired yet. - raise NotImplementedError( - "server_version is not yet implemented for the Rust Postgres backend" - ) + """Returns a string giving the server version. For example: '14.4'.""" + numver = self._version + assert numver is not None, "check_database must be called first" + # Supported versions are all >= 10, so use the two-part form. + # https://www.postgresql.org/docs/current/libpq-status.html#LIBPQ-PQSERVERVERSION + return "%i.%i" % (numver / 10000, numver % 10000) diff --git a/synapse/storage/rust_dbapi.py b/synapse/storage/rust_dbapi.py index 3aa696d722..78486cea39 100644 --- a/synapse/storage/rust_dbapi.py +++ b/synapse/storage/rust_dbapi.py @@ -32,11 +32,159 @@ routing change in ``LoggingTransaction`` to reach a shim-backed implementation; that is a separate follow-up. """ -from typing import TYPE_CHECKING, Any, Iterator, Sequence +import logging +from typing import TYPE_CHECKING, Any, Iterator, Mapping, Sequence + +from synapse.synapse_rust.database import postgres if TYPE_CHECKING: from synapse.storage.types import SQLQueryParameters +logger = logging.getLogger(__name__) + + +def _quote_dsn_value(value: Any) -> str: + """Quote a value for a libpq keyword/value connection string. + + Values with spaces or quotes must be single-quoted with `\\` and `'` + backslash-escaped; an empty value must be `''`. + """ + s = str(value) + if s == "" or any(c in s for c in " '\\"): + escaped = s.replace("\\", "\\\\").replace("'", "\\'") + return f"'{escaped}'" + return s + + +# psycopg2 accepts a few connection kwargs that are *not* libpq keywords (or +# not tokio_postgres's spelling of them) and translates them itself. The Rust +# pool instead parses a strict libpq DSN (via ``tokio_postgres``), which only +# knows its own keywords, so we map the aliases here. ``database`` -> +# ``dbname`` is the important one: Synapse's sample config and most real +# deployments spell the database name ``database``. ``keepalives_count`` is +# libpq's name for what tokio_postgres calls ``keepalives_retries`` (and is +# what docs/postgres.md recommends setting). +_PSYCOPG2_KEY_ALIASES = { + "database": "dbname", + "keepalives_count": "keepalives_retries", +} + +# The keywords tokio_postgres's DSN parser accepts (see ``Config::param`` in +# tokio-postgres 0.7.x, pinned by Cargo.lock; parity is asserted by +# ``test_supported_dsn_keys_are_accepted_by_the_parser``). It hard-errors on +# anything else, unlike libpq/psycopg2 which accept a much wider set — so only +# these keys may reach the DSN. ``sslmode`` is deliberately absent: the +# ``ssl*`` keys are split into TLS params before the DSN is built (see +# ``split_ssl_params``), and a DSN-level ``sslmode`` would be silently +# overridden by the pool's TLS setup — better to fail loudly. +_SUPPORTED_DSN_KEYS = frozenset( + { + "application_name", + "channel_binding", + "connect_timeout", + "dbname", + "host", + "hostaddr", + "keepalives", + "keepalives_idle", + "keepalives_interval", + "keepalives_retries", + "load_balance_hosts", + "options", + "password", + "port", + "sslnegotiation", + "target_session_attrs", + "tcp_user_timeout", + "user", + } +) + +# libpq keywords whose absence cannot change where we connect, how we +# authenticate, or whether the connection is encrypted; these are dropped with +# a warning, for compatibility with psycopg2-era configs. Anything else +# unknown is a hard error: silently dropping e.g. ``service`` or ``sslcrl`` +# could connect to the wrong database or downgrade security. +_DROPPABLE_DSN_KEYS = frozenset( + { + # tokio_postgres always talks UTF8 (and Synapse requires a UTF8 DB). + "client_encoding", + # Deprecated in libpq and a no-op since PostgreSQL 14. + "sslcompression", + } +) + + +def build_dsn(params: Mapping[str, Any]) -> str: + """Build a libpq keyword/value DSN from psycopg2-style connection kwargs. + + Synapse's database `args` are psycopg2 connection kwargs — mostly libpq + keywords (``user``, ``host``, ``port``, ``password``, …), but ``database`` + is a psycopg2 alias for libpq's ``dbname`` (see ``_PSYCOPG2_KEY_ALIASES``). + The Rust pool takes a strict libpq DSN string rather than kwargs, so join + them into one, mapping any aliases to their real keyword. ``None`` values + are skipped, matching psycopg2's handling of `None` kwargs (fall back to the + libpq default). + + Keywords tokio_postgres doesn't know cannot go into the DSN (its parser + hard-errors with no hint about which key is at fault). The known-harmless + ones (``_DROPPABLE_DSN_KEYS``) are dropped with a warning; anything else + raises, because silently dropping a key like ``service`` or ``passfile`` + could change *which database we connect to* or quietly downgrade security. + + Raises: + ValueError: for an arg the Rust driver can't honour, or when an alias + and its target are both set (e.g. ``database`` and ``dbname``), + which psycopg2 also rejects. + """ + parts = [] + for key, value in params.items(): + if value is None: + continue + mapped = _PSYCOPG2_KEY_ALIASES.get(key, key) + if mapped != key and params.get(mapped) is not None: + raise ValueError( + f"database config args set both {key!r} and {mapped!r}; " + "remove one of them" + ) + if mapped not in _SUPPORTED_DSN_KEYS: + if mapped in _DROPPABLE_DSN_KEYS: + logger.warning( + "Ignoring database connection argument %r: " + "not supported by the native Rust driver", + key, + ) + continue + raise ValueError( + f"database config arg {key!r} is not supported by the native " + "Rust driver; remove it from `args` or disable " + "`use_rust_driver`" + ) + parts.append(f"{mapped}={_quote_dsn_value(value)}") + return " ".join(parts) + + +def connect( + dsn: str, + *, + synchronous_commit: bool = True, + statement_timeout_ms: int | None = None, +) -> "Connection": + """Open a single standalone connection for bootstrap/one-off use. + + The Rust shim is pool-only, so a lone connection is a pool of one; the + returned :class:`Connection` keeps that pool alive for its lifetime. Used by + ``make_conn`` for the startup connection that runs schema preparation before + the real pool exists. + """ + pool = postgres.ConnectionPool( + dsn, + 1, + synchronous_commit=synchronous_commit, + statement_timeout_ms=statement_timeout_ms, + ) + return Connection(pool.connect(), pool=pool) + class Cursor: """A DBAPI2 cursor wrapping a Rust shim cursor.""" @@ -120,8 +268,12 @@ class Connection: engine-facing methods delegate straight to the shim. """ - def __init__(self, conn: Any) -> None: + def __init__(self, conn: Any, pool: Any = None) -> None: self._conn = conn + # A pool this connection owns (a bootstrap "pool of one"), kept alive for + # the connection's lifetime and closed with it. `None` for connections + # borrowed from a shared pool. + self._pool = pool def cursor(self) -> Cursor: return Cursor(self._conn.cursor()) @@ -134,6 +286,8 @@ class Connection: def close(self) -> None: self._conn.close() + if self._pool is not None: + self._pool.close() # -- engine-facing methods (see RustPostgresEngine) --------------------- diff --git a/tests/storage/test_rust_dbapi.py b/tests/storage/test_rust_dbapi.py index 812e81022f..fc72455aa5 100644 --- a/tests/storage/test_rust_dbapi.py +++ b/tests/storage/test_rust_dbapi.py @@ -14,8 +14,9 @@ (:mod:`synapse.storage.rust_dbapi`), including driving a real ``LoggingTransaction`` through it.""" +from synapse.config.database import DatabaseConnectionConfig from synapse.storage import rust_dbapi -from synapse.storage.database import LoggingDatabaseConnection +from synapse.storage.database import LoggingDatabaseConnection, make_conn from synapse.storage.engines.postgres_rust import RustPostgresEngine from synapse.synapse_rust.database import postgres @@ -46,6 +47,129 @@ def _build_dsn() -> str: return " ".join(parts) +class BuildDsnTestCase(unittest.TestCase): + """`build_dsn` turns psycopg2-style kwargs into a libpq DSN (no database).""" + + def test_joins_keywords(self) -> None: + self.assertEqual( + rust_dbapi.build_dsn( + {"dbname": "synapse", "user": "u", "host": "db", "port": 5432} + ), + "dbname=synapse user=u host=db port=5432", + ) + + def test_quotes_values_needing_it(self) -> None: + # Spaces / quotes / backslashes get single-quoted and escaped; empty → ''. + self.assertEqual( + rust_dbapi.build_dsn({"password": "p a'ss\\x", "options": ""}), + "password='p a\\'ss\\\\x' options=''", + ) + + def test_maps_keepalives_count_to_keepalives_retries(self) -> None: + # libpq (and docs/postgres.md's example config) spell it + # `keepalives_count`; tokio_postgres spells it `keepalives_retries`. + self.assertEqual( + rust_dbapi.build_dsn({"keepalives": 1, "keepalives_count": 3}), + "keepalives=1 keepalives_retries=3", + ) + + def test_alias_colliding_with_its_target_raises(self) -> None: + # Both spellings set: psycopg2 rejects database+dbname with a + # TypeError; silently letting one win would be config-order lottery. + with self.assertRaises(ValueError): + rust_dbapi.build_dsn({"database": "a", "dbname": "b"}) + with self.assertRaises(ValueError): + rust_dbapi.build_dsn({"keepalives_count": 3, "keepalives_retries": 5}) + + def test_drops_known_harmless_keywords_with_a_warning(self) -> None: + # These libpq keys can't change the connection target, auth, or + # security posture, so psycopg2-era configs carrying them keep working. + self.assertEqual( + rust_dbapi.build_dsn( + {"dbname": "d", "client_encoding": "UTF8", "sslcompression": 0} + ), + "dbname=d", + ) + + def test_rejects_keywords_that_could_change_target_or_security(self) -> None: + # Silently dropping these would connect to the wrong database + # (service/passfile) or downgrade security (sslcrl, gssencmode, ...): + # fail loudly instead. + for key, value in ( + ("service", "synapse-prod"), + ("passfile", "/etc/pgpass"), + ("sslcrl", "/etc/crl.pem"), + ("gssencmode", "require"), + ("ssl_min_protocol_version", "TLSv1.3"), + ): + with self.assertRaises(ValueError, msg=key) as ctx: + rust_dbapi.build_dsn({"dbname": "d", key: value}) + self.assertIn(key, str(ctx.exception)) + + def test_supported_dsn_keys_are_accepted_by_the_parser(self) -> None: + # _SUPPORTED_DSN_KEYS mirrors tokio_postgres's Config::param keyword + # set; if the crate is upgraded and a key is renamed or removed, this + # catches the drift (the pool parses its DSN eagerly, no server + # needed). + samples = { + "channel_binding": "disable", + "connect_timeout": "5", + "hostaddr": "127.0.0.1", + "keepalives": "1", + "load_balance_hosts": "disable", + "sslnegotiation": "postgres", + "target_session_attrs": "any", + } + for key in sorted(rust_dbapi._SUPPORTED_DSN_KEYS): + value = samples.get(key, "1" if key.startswith("keepalives") else "x") + if key in ("port", "tcp_user_timeout"): + value = "5432" + pool = postgres.ConnectionPool(f"host=h {key}={value}") + pool.close() + + +@skip_unless( + bool(USE_POSTGRES_FOR_TESTS), "requires a Postgres server (set SYNAPSE_POSTGRES)" +) +class RustStartupTestCase(unittest.TestCase): + """The startup path: `make_conn` + `check_database` for the Rust engine.""" + + def _db_config(self) -> DatabaseConnectionConfig: + args: dict = {"dbname": POSTGRES_BASE_DB} + if POSTGRES_USER is not None: + args["user"] = POSTGRES_USER + if POSTGRES_HOST is not None: + args["host"] = POSTGRES_HOST + if POSTGRES_PORT is not None: + args["port"] = POSTGRES_PORT + if POSTGRES_PASSWORD is not None: + args["password"] = POSTGRES_PASSWORD + # `name` must be a recognised engine; the Rust engine is selected by + # passing a RustPostgresEngine to make_conn, not by the config name. + return DatabaseConnectionConfig("master", {"name": "psycopg2", "args": args}) + + def test_make_conn_check_database_and_query(self) -> None: + engine = RustPostgresEngine({}) + db_conn = make_conn( + db_config=self._db_config(), + engine=engine, + default_txn_name="startup", + server_name="test", + ) + try: + # A bootstrap connection the engine can validate over a cursor. + engine.check_database(db_conn) + self.assertRegex(engine.server_version, r"^\d+\.\d+$") + + # And it's a working connection. + with db_conn.cursor(txn_name="startup") as cur: + cur.execute("SELECT 1") + self.assertEqual(cur.fetchone(), (1,)) + db_conn.commit() + finally: + db_conn.close() + + @skip_unless( bool(USE_POSTGRES_FOR_TESTS), "requires a Postgres server (set SYNAPSE_POSTGRES)" )