Wire the startup path for the Rust Postgres backend

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

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,

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)

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

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)"
 )