Map Postgres errors onto a DBAPI2 exception hierarchy

Previously every `tokio_postgres` error became a bare `RuntimeError`. Synapse's
transaction driver, though, branches on the *type* of a database error and on
its `pgcode`: `new_transaction` retries `OperationalError`, retries deadlocks it
recognises via `is_deadlock` (which reads `pgcode`) on a `DatabaseError`, and
`simple_upsert` retries `IntegrityError`. With everything collapsed to
`RuntimeError` none of that fired.

Add just the distinctions Synapse acts on, rather than psycopg2's full PEP-249
hierarchy: `Error` -> `DatabaseError` -> {`OperationalError`, `IntegrityError`},
exposed on the `postgres` submodule, each instance tagged with `pgcode` (the
SQLSTATE string, or `None`). A small classifier maps the SQLSTATE class:
constraint violations (`23`) to `IntegrityError`, connection/resource classes
(`08`/`53`/`57`/`58`) to `OperationalError`, everything else (incl. `40*`
deadlocks, which retry via `pgcode`) to `DatabaseError`. Codeless errors are
split with `is_closed()`: a lost connection is operational, any other (a bad
parameter, a failed connect) is a plain `DatabaseError` so it isn't retried.

Errors surfacing while a result stream is drained (the usual case for an
`INSERT` constraint violation) now route through the same mapping, so they carry
the right class and `pgcode` too.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

rust/src/database/postgres/cursor_state.rs

@@ -37,6 +37,7 @@ use pyo3::{
 use tokio_postgres::RowStream;
 
 use crate::database::postgres::{
+    errors::pg_err_to_py,
     helpers::{BlockingPostgres, BlockingPostgresStream as _},
     value::pg_row_to_py,
 };
@@ -83,7 +84,11 @@ impl CursorRowStream for RowStream {
     }
 
     fn stream_err(err: &Self::Error) -> PyErr {
-        PyRuntimeError::new_err(format!("error fetching row from postgres: {err}"))
+        // Route through the shared mapping so a server error surfacing while
+        // the result stream is drained (e.g. a constraint violation on an
+        // `INSERT`) becomes the right DBAPI2 exception with its `pgcode`, just
+        // as it would if it surfaced at `execute` time.
+        pg_err_to_py(err)
     }
 }
 

rust/src/database/postgres/errors.rs

@@ -0,0 +1,281 @@
+//! The DBAPI2 exception hierarchy for the Postgres backend, and the mapping
+//! from a [`tokio_postgres`] error onto it.
+//!
+//! Synapse's transaction driver (`synapse.storage.database.new_transaction`)
+//! and the Postgres engine branch on the *type* of the exception a database
+//! call raises, and on its `pgcode` (the 5-character SQLSTATE). To be a drop-in
+//! for psycopg2 we therefore need to raise exceptions of the right class and
+//! carry a `pgcode`.
+//!
+//! Rather than reproduce psycopg2's full PEP-249 hierarchy (ten classes) and
+//! its complete SQLSTATE→class table, we expose only the distinctions Synapse
+//! actually acts on:
+//!
+//! * [`Error`] — the base every database error derives from. Synapse catches it
+//!   when a rollback itself fails.
+//! * [`DatabaseError`] — a server-side error. Synapse catches this and calls
+//!   `is_deadlock`, which reads `pgcode` to spot serialization/deadlock failures
+//!   (`40001`/`40P01`) and retry them.
+//! * [`OperationalError`] — a transient/connection-level failure ("the database
+//!   disappeared mid-transaction"). Synapse catches this and retries.
+//! * [`IntegrityError`] — a constraint violation. Synapse catches this to retry
+//!   upserts and to handle insert races.
+//! * [`ProgrammingError`] — a SQL-level mistake (syntax error, duplicate table
+//!   or index, …; SQLSTATE class `42`). Caught by the search store's GIN-index
+//!   migration to ignore "already exists".
+//!
+//! Every other Postgres error (data errors, …) surfaces as a plain
+//! [`DatabaseError`]. Nothing in Synapse catches the remaining psycopg2 classes
+//! (`DataError`, `InternalError`, …), so collapsing them is invisible at
+//! runtime. If full [`DBAPI2Module`] conformance is needed later (when a Rust
+//! engine is wired up) the remaining PEP-249 names can be added as aliases of
+//! [`DatabaseError`].
+//!
+//! [`DBAPI2Module`]: (see `synapse/storage/types.py`)
+
+use pyo3::exceptions::PyException;
+use pyo3::prelude::*;
+use pyo3::{create_exception, types::PyModule};
+
+create_exception!(
+    postgres,
+    Error,
+    PyException,
+    "Base class for every error raised by the Rust Postgres backend (PEP-249 `Error`)."
+);
+create_exception!(
+    postgres,
+    DatabaseError,
+    Error,
+    "A server-side database error. Carries the SQLSTATE as `pgcode`."
+);
+create_exception!(
+    postgres,
+    OperationalError,
+    DatabaseError,
+    "A transient/connection-level failure that is worth retrying."
+);
+create_exception!(
+    postgres,
+    IntegrityError,
+    DatabaseError,
+    "A constraint violation (e.g. a unique or foreign-key violation)."
+);
+create_exception!(
+    postgres,
+    ProgrammingError,
+    DatabaseError,
+    "A SQL-level mistake (syntax error, duplicate table/index, undefined column, ...)."
+);
+
+/// Build the Python exception for a Postgres failure, tagging it with `pgcode`.
+///
+/// `code` is the SQLSTATE (`None` for an error that never got a server
+/// response). When present, the class is chosen from the SQLSTATE *class* (its
+/// first two characters):
+///
+/// * `23` (integrity constraint violation) → [`IntegrityError`]
+/// * `08`/`53`/`57`/`58` (connection, resource, operator-intervention,
+///   system errors) → [`OperationalError`]
+/// * `42` (syntax error or access rule violation, e.g. `42P07` duplicate
+///   table/index) → [`ProgrammingError`], matching psycopg2's mapping
+/// * everything else → [`DatabaseError`]
+///
+/// Note deadlock/serialization failures (`40001`/`40P01`) deliberately fall
+/// into the [`DatabaseError`] bucket, not [`OperationalError`]: Synapse retries
+/// them via `is_deadlock`, which only needs a `DatabaseError` with the right
+/// `pgcode`.
+///
+/// A codeless error is one that never reached, or never heard back from, the
+/// server. We can't inspect its kind, but `closed` (from
+/// [`tokio_postgres::Error::is_closed`]) tells us whether it was a lost
+/// connection: if so it's [`OperationalError`] and worth retrying; otherwise
+/// it's a client-side problem (a bad parameter, a failed connect) that
+/// shouldn't be retried, so it becomes a plain [`DatabaseError`].
+fn new_err_for(py: Python<'_>, code: Option<&str>, closed: bool, msg: &str) -> PyErr {
+    let err = match code.map(sqlstate_class) {
+        Some("23") => IntegrityError::new_err(msg.to_string()),
+        Some("08" | "53" | "57" | "58") => OperationalError::new_err(msg.to_string()),
+        Some("42") => ProgrammingError::new_err(msg.to_string()),
+        // A server error we don't single out, e.g. a deadlock (`40*`) or a
+        // data error (`22*`), …
+        Some(_) => DatabaseError::new_err(msg.to_string()),
+        // No SQLSTATE: a lost connection is operational (retry); any other
+        // codeless error is client-side and propagates as a `DatabaseError`.
+        None if closed => OperationalError::new_err(msg.to_string()),
+        None => DatabaseError::new_err(msg.to_string()),
+    };
+
+    // Attach `pgcode` (the SQLSTATE string, or Python `None`) so engine code
+    // such as `is_deadlock` can read `error.pgcode` exactly as it does for
+    // psycopg2. Setting an attribute on a fresh exception instance does not
+    // fail in practice, so a failure here is not worth propagating.
+    let _ = err.value(py).setattr("pgcode", code);
+
+    err
+}
+
+/// The SQLSTATE *class*: the first two characters of a 5-character code.
+fn sqlstate_class(code: &str) -> &str {
+    &code[..2.min(code.len())]
+}
+
+/// Map a [`tokio_postgres`] error into the appropriate Python exception.
+///
+/// Takes the error by reference so it can be called both from the `block_on`
+/// helpers (which own the error) and from the cursor's row-stream draining
+/// (which only has a borrow).
+///
+/// For a server-reported error the message includes the SQLSTATE, the server's
+/// message and, when present, its DETAIL and HINT — `tokio_postgres`'s own
+/// `Display` is just "db error", which buries the reason (psycopg2 surfaces
+/// the full server message).
+pub(crate) fn pg_err_to_py(e: &tokio_postgres::Error) -> PyErr {
+    let code = e.code().map(|c| c.code());
+    let msg = match e.as_db_error() {
+        Some(db_err) => {
+            let mut msg = format!(
+                "postgres error {}: {}",
+                db_err.code().code(),
+                db_err.message()
+            );
+            if let Some(detail) = db_err.detail() {
+                msg.push_str(&format!(" DETAIL: {detail}"));
+            }
+            if let Some(hint) = db_err.hint() {
+                msg.push_str(&format!(" HINT: {hint}"));
+            }
+            msg
+        }
+        None => format!("postgres error: {e}"),
+    };
+    Python::attach(|py| new_err_for(py, code, e.is_closed(), &msg))
+}
+
+/// Add the exception classes to the `postgres` submodule so the module conforms
+/// to (the load-bearing part of) Synapse's `DBAPI2Module` protocol.
+pub(crate) fn register_exceptions(py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> {
+    m.add("Error", py.get_type::<Error>())?;
+    m.add("DatabaseError", py.get_type::<DatabaseError>())?;
+    m.add("OperationalError", py.get_type::<OperationalError>())?;
+    m.add("IntegrityError", py.get_type::<IntegrityError>())?;
+    m.add("ProgrammingError", py.get_type::<ProgrammingError>())?;
+    Ok(())
+}
+
+#[cfg(test)]
+mod tests {
+    //! These tests don't touch Postgres: `new_err_for` takes the SQLSTATE as a
+    //! plain string, so the whole classification and `pgcode` tagging can be
+    //! exercised without a live database (a `tokio_postgres::Error` can't be
+    //! constructed by hand anyway).
+
+    use pyo3::PyTypeInfo;
+
+    use super::*;
+
+    /// Assert `code` maps to exception type `T`, is a `DatabaseError` and an
+    /// `Error`, and carries the expected `pgcode`.
+    fn assert_maps<T: PyTypeInfo>(code: &str) {
+        Python::attach(|py| {
+            let err = new_err_for(py, Some(code), false, "boom");
+            let value = err.value(py);
+            assert!(
+                value.is_instance_of::<T>(),
+                "{code} did not map to the expected type"
+            );
+            assert!(value.is_instance_of::<DatabaseError>());
+            assert!(value.is_instance_of::<Error>());
+
+            let pgcode: String = value.getattr("pgcode").unwrap().extract().unwrap();
+            assert_eq!(pgcode, code);
+        });
+    }
+
+    #[test]
+    fn integrity_class_maps_to_integrity_error() {
+        Python::initialize();
+        // Unique violation, foreign-key violation, not-null violation.
+        assert_maps::<IntegrityError>("23505");
+        assert_maps::<IntegrityError>("23503");
+    }
+
+    #[test]
+    fn connection_and_resource_classes_map_to_operational_error() {
+        Python::initialize();
+        assert_maps::<OperationalError>("08006"); // connection failure
+        assert_maps::<OperationalError>("53100"); // disk full
+        assert_maps::<OperationalError>("57014"); // query canceled
+        assert_maps::<OperationalError>("58000"); // system error
+    }
+
+    #[test]
+    fn programming_class_maps_to_programming_error() {
+        Python::initialize();
+        // Syntax error, duplicate table/index (what the GIN-index migration
+        // catches to ignore "already exists").
+        assert_maps::<ProgrammingError>("42601");
+        assert_maps::<ProgrammingError>("42P07");
+    }
+
+    #[test]
+    fn other_server_errors_map_to_plain_database_error() {
+        Python::initialize();
+        // A data error is a `DatabaseError` but none of the specific classes.
+        Python::attach(|py| {
+            let value = new_err_for(py, Some("22012"), false, "boom");
+            let value = value.value(py);
+            assert!(value.is_instance_of::<DatabaseError>());
+            assert!(!value.is_instance_of::<OperationalError>());
+            assert!(!value.is_instance_of::<IntegrityError>());
+            assert!(!value.is_instance_of::<ProgrammingError>());
+        });
+    }
+
+    #[test]
+    fn deadlock_and_serialization_are_retryable_database_errors() {
+        Python::initialize();
+        // The behaviour Synapse's `is_deadlock` relies on: a `DatabaseError`
+        // (so the `isinstance` check passes) carrying the right `pgcode`.
+        for code in ["40001", "40P01"] {
+            Python::attach(|py| {
+                let err = new_err_for(py, Some(code), false, "boom");
+                let value = err.value(py);
+                assert!(value.is_instance_of::<DatabaseError>());
+                let pgcode: String = value.getattr("pgcode").unwrap().extract().unwrap();
+                assert_eq!(pgcode, code);
+            });
+        }
+    }
+
+    #[test]
+    fn closed_connection_error_is_operational() {
+        Python::initialize();
+        Python::attach(|py| {
+            // A lost connection (codeless, `is_closed()`) is operational, so
+            // Synapse retries it.
+            let err = new_err_for(py, None, true, "connection closed");
+            let value = err.value(py);
+            assert!(value.is_instance_of::<OperationalError>());
+            assert!(value.is_instance_of::<DatabaseError>());
+            // `pgcode` is present but `None`, so `error.pgcode in (...)` is a
+            // safe membership test rather than an `AttributeError`.
+            assert!(value.getattr("pgcode").unwrap().is_none());
+        });
+    }
+
+    #[test]
+    fn other_codeless_error_is_a_plain_database_error() {
+        Python::initialize();
+        Python::attach(|py| {
+            // A client-side error that isn't a lost connection (a bad
+            // parameter, a failed connect) is *not* retryable, so it is a plain
+            // `DatabaseError`, not `OperationalError`.
+            let err = new_err_for(py, None, false, "error serializing parameter");
+            let value = err.value(py);
+            assert!(value.is_instance_of::<DatabaseError>());
+            assert!(!value.is_instance_of::<OperationalError>());
+            assert!(value.getattr("pgcode").unwrap().is_none());
+        });
+    }
+}

rust/src/database/postgres/helpers.rs

@@ -29,7 +29,7 @@ use std::{future::Future, pin::Pin};
 use futures::{stream::Fuse, FutureExt, StreamExt};
 use pyo3::{marker::Ungil, PyResult, Python};
 
-use crate::database::{postgres::pg_err_to_py, runtime::runtime};
+use crate::database::{postgres::errors::pg_err_to_py, runtime::runtime};
 
 /// Block on a future on the shared runtime, releasing the GIL while we wait.
 pub trait BlockingPostgres
@@ -53,7 +53,7 @@ where
 {
     /// Block on `self` and convert a Postgres error into a `PyErr`.
     fn block_on_result(self, py: Python<'_>) -> PyResult<T> {
-        self.block_on(py).map_err(pg_err_to_py)
+        self.block_on(py).map_err(|e| pg_err_to_py(&e))
     }
 }
 

rust/src/database/postgres/mod.rs

@@ -14,18 +14,21 @@ use crate::database::runtime::runtime;
 
 mod connection;
 mod cursor_state;
+mod errors;
 mod helpers;
 mod libpq;
 mod value;
 
-/// Register the `postgres` submodule (the `Connection` / `Cursor` classes and
-/// the `connect` factory) under the parent `database` module.
+/// Register the `postgres` submodule (the `Connection` / `Cursor` classes, the
+/// DBAPI2 exception hierarchy and the `connect` factory) under the parent
+/// `database` module.
 pub fn register_module(py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> {
     let child = PyModule::new(py, "postgres")?;
 
     child.add_class::<connection::Connection>()?;
     child.add_class::<connection::Cursor>()?;
     child.add_function(wrap_pyfunction!(connect, &child)?)?;
+    errors::register_exceptions(py, &child)?;
 
     m.add_submodule(&child)?;
 
@@ -38,11 +41,6 @@ pub fn register_module(py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()>
     Ok(())
 }
 
-/// Map a [`tokio_postgres`] error into a Python `RuntimeError`.
-fn pg_err_to_py(e: tokio_postgres::Error) -> PyErr {
-    PyRuntimeError::new_err(format!("postgres error: {e}"))
-}
-
 /// Open a new Postgres connection from a libpq-style DSN.
 ///
 /// Blocks until the connection is established, then spawns the long-lived

synapse/storage/databases/main/search.py

@@ -286,9 +286,9 @@ class SearchBackgroundUpdateStore(SearchWorkerStore):
                 # if we skipped the conversion to GIST, we may already/still
                 # have an event_search_fts_idx; unfortunately postgres 9.4
                 # doesn't support CREATE INDEX IF EXISTS so we just catch the
-                # exception and ignore it.
-                import psycopg2
-
+                # exception and ignore it. Catch the error class via the
+                # engine's module so this works on both the psycopg2 and the
+                # native Rust drivers.
                 try:
                     c.execute(
                         """
@@ -296,7 +296,7 @@ class SearchBackgroundUpdateStore(SearchWorkerStore):
                         ON event_search USING GIN (vector)
                         """
                     )
-                except psycopg2.ProgrammingError as e:
+                except self.database_engine.module.ProgrammingError as e:
                     logger.warning(
                         "Ignoring error %r when trying to switch from GIST to GIN", e
                     )

tests/synapse_rust/test_database_postgres.py

@@ -97,9 +97,9 @@ class PostgresConnectionTestCase(unittest.TestCase):
     # ------------------------------------------------------------------
 
     def test_connect_bad_dsn_raises(self) -> None:
-        # A syntactically valid but unconnectable DSN should raise rather than
-        # return a half-open connection.
-        with self.assertRaises(RuntimeError):
+        # A syntactically valid but unconnectable DSN should raise one of our
+        # DBAPI2 errors rather than return a half-open connection.
+        with self.assertRaises(postgres.Error):
             postgres.connect("host=127.0.0.1 port=1 dbname=does_not_exist")
 
     # ------------------------------------------------------------------
@@ -455,8 +455,12 @@ class PostgresConnectionTestCase(unittest.TestCase):
             cursor.execute("CREATE TEMP TABLE oor (x int4)")
             cursor.execute("INSERT INTO oor VALUES ($1)", [10**12])
 
-        with self.assertRaises(RuntimeError):
+        # This is a client-side encoding error (no SQLSTATE), so it surfaces as
+        # a plain DatabaseError -- not an OperationalError, since retrying a
+        # deterministic bad value would be pointless.
+        with self.assertRaises(postgres.DatabaseError) as ctx:
             run_interaction(self.conn, interaction)
+        self.assertNotIsInstance(ctx.exception, postgres.OperationalError)
 
     # ------------------------------------------------------------------
     # rowcount()
@@ -570,10 +574,10 @@ class PostgresConnectionTestCase(unittest.TestCase):
 
         def bad_sql(cursor: Any) -> None:
             # A syntax error: the server rejects this during prepare, which
-            # surfaces as a RuntimeError and rolls the transaction back.
+            # surfaces as a DatabaseError and rolls the transaction back.
             cursor.execute("SELECT FROM WHERE not valid sql")
 
-        with self.assertRaises(RuntimeError):
+        with self.assertRaises(postgres.DatabaseError):
             run_interaction(self.conn, bad_sql)
 
         # The transaction was rolled back and the connection handed back clean,
@@ -593,12 +597,15 @@ class PostgresConnectionTestCase(unittest.TestCase):
             cursor.execute("INSERT INTO uniq VALUES (1)")
             # Duplicate key. The error is reported by the server while the
             # statement's result stream is driven, so we drain it (via
-            # rowcount) to surface it as a RuntimeError.
+            # rowcount) to surface it -- as an IntegrityError, the same class
+            # (with the same 23505 pgcode) it would carry had it surfaced at
+            # execute time.
             cursor.execute("INSERT INTO uniq VALUES (1)")
             cursor.rowcount()
 
-        with self.assertRaises(RuntimeError):
+        with self.assertRaises(postgres.IntegrityError) as ctx:
             run_interaction(self.conn, violate)
+        self.assertEqual(ctx.exception.pgcode, "23505")
 
         # TEMP table lived only in the rolled-back transaction; the connection
         # itself is fine.
@@ -810,6 +817,71 @@ class PostgresConnectionDrivenTestCase(unittest.TestCase):
         self.conn.rollback()
 
 
+@unittest.skip_unless(
+    bool(USE_POSTGRES_FOR_TESTS), "requires a Postgres server (set SYNAPSE_POSTGRES)"
+)
+class PostgresErrorMappingTestCase(unittest.TestCase):
+    """The DBAPI2 exception hierarchy and the SQLSTATE→exception mapping.
+
+    Synapse's transaction driver branches on the *type* of the exception a
+    database call raises (``OperationalError`` → retry, ``IntegrityError`` →
+    retry upserts) and on its ``pgcode`` (``is_deadlock``). These tests check
+    the Rust backend raises the right class and carries a ``pgcode``, the way
+    psycopg2 does.
+    """
+
+    def setUp(self) -> None:
+        self.conn = postgres.connect(_build_dsn())
+
+    def tearDown(self) -> None:
+        del self.conn
+
+    def _exec_commit(self, sql: str) -> None:
+        """Run a single statement and commit it (its own transaction)."""
+        self.conn.cursor().execute(sql)
+        self.conn.commit()
+
+    # -- the hierarchy exposed on the module --------------------------------
+
+    def test_module_exposes_dbapi2_hierarchy(self) -> None:
+        """The exception attributes Synapse's engine code and DBAPI2Module
+        protocol rely on, with the expected subclass links."""
+        self.assertTrue(issubclass(postgres.DatabaseError, postgres.Error))
+        self.assertTrue(issubclass(postgres.OperationalError, postgres.DatabaseError))
+        self.assertTrue(issubclass(postgres.IntegrityError, postgres.DatabaseError))
+
+    # -- SQLSTATE → exception class -----------------------------------------
+
+    def test_unique_violation_is_integrity_error(self) -> None:
+        """A constraint violation raises ``IntegrityError`` with pgcode 23505."""
+        table = "rust_pg_err_integrity"
+        try:
+            self._exec_commit(f"CREATE TABLE {table} (id int PRIMARY KEY)")
+            self._exec_commit(f"INSERT INTO {table} VALUES (1)")
+
+            cursor = self.conn.cursor()
+            with self.assertRaises(postgres.IntegrityError) as ctx:
+                cursor.execute(f"INSERT INTO {table} VALUES (1)")
+                # The INSERT's error is reported while its result stream is
+                # driven, so drain it (via rowcount) to surface it.
+                cursor.rowcount()
+            self.assertEqual(ctx.exception.pgcode, "23505")
+            self.conn.rollback()
+        finally:
+            self._exec_commit(f"DROP TABLE IF EXISTS {table}")
+
+    def test_undefined_table_is_plain_database_error(self) -> None:
+        """An error we don't single out surfaces as a plain ``DatabaseError``
+        (not one of the specialised subclasses), still carrying its pgcode."""
+        cursor = self.conn.cursor()
+        with self.assertRaises(postgres.DatabaseError) as ctx:
+            cursor.execute("SELECT * FROM rust_pg_no_such_table")
+        self.assertNotIsInstance(ctx.exception, postgres.OperationalError)
+        self.assertNotIsInstance(ctx.exception, postgres.IntegrityError)
+        self.assertEqual(ctx.exception.pgcode, "42P01")  # undefined_table
+        self.conn.rollback()
+
+
 @unittest.skip_unless(
     bool(USE_POSTGRES_FOR_TESTS) and POSTGRES_HOST in (None, "", "localhost"),
     "requires Postgres reachable on libpq's default host",