shared/database.py

"""Short server-only PostgreSQL transactions with bounded failure and TLS."""

from collections.abc import Iterator
from contextlib import contextmanager
from typing import Any


class StorageUnavailable(RuntimeError):
    """Storage cannot authorize work; callers must not silently fall back."""


@contextmanager
def transaction(dsn: str, *, lock: int | None = None) -> Iterator[Any]:
    try:
        import psycopg
        from psycopg.conninfo import conninfo_to_dict
        from psycopg.rows import dict_row
    except ImportError:
        raise StorageUnavailable(
            "Install the storage extra before enabling the database."
        ) from None
    try:
        params = conninfo_to_dict(dsn)
        if params.get("host") not in {"localhost", "127.0.0.1", "::1"} and params.get(
            "sslmode"
        ) not in {"require", "verify-ca", "verify-full"}:
            raise StorageUnavailable("The remote database connection requires TLS.")
        with psycopg.connect(
            dsn, connect_timeout=5, prepare_threshold=None, row_factory=dict_row
        ) as conn:
            conn.execute("SET LOCAL lock_timeout = '5s'")
            conn.execute("SET LOCAL statement_timeout = '5s'")
            if lock is not None:
                conn.execute("SELECT pg_advisory_xact_lock(%s)", (lock,))
            yield conn
    except psycopg.Error:
        # Connection errors may contain passwords/DSNs. Do not retain their cause.
        raise StorageUnavailable("Database storage is temporarily unavailable.") from None