gnomon-postgres 1.0.0__py3-none-any.whl

gnomon/asyncpg.py

@@ -0,0 +1,171 @@
+"""asyncpg for the async SDK.
+
+    import asyncpg
+    from gnomon import AsyncGnomon
+    from gnomon.asyncpg import asyncpg_source
+
+    app = await asyncpg.create_pool("postgresql://app@db/shop")
+    runs = await asyncpg.create_pool("postgresql://runs@db/shop")
+    gn = AsyncGnomon(app=asyncpg_source(app), runs=asyncpg_source(runs))
+
+The SDK speaks psycopg's interface; this hands it connections from an asyncpg pool (or a single
+asyncpg connection) that answer the part of it the SDK uses: ``%s`` placeholders become ``$n``,
+rows come back as dicts (or tuples), row counts come from the command status, and PostgreSQL's
+errors come back as psycopg errors carrying the SQLSTATE and hint, so they map to the SDK's typed
+errors. A statement outside a transaction runs in one of its own, so behind a transaction pooler
+(pgbouncer) its prepare and execute stay on one server connection.
+"""
+from __future__ import annotations
+
+import asyncio
+import contextlib
+import re
+from typing import Any, AsyncIterator, Optional
+
+import psycopg
+
+
+def to_dollar(sql: str) -> str:
+    """psycopg's placeholders (%s, %% for a percent sign) as PostgreSQL's ($1, $2, ...), leaving
+    quoted strings, identifiers, dollar-quoted bodies and comments alone."""
+    out, i, n, k = [], 0, len(sql), 0
+    while i < n:
+        c = sql[i]
+        if c in "'\"":
+            j = i + 1
+            while j < n:
+                if sql[j] == c:
+                    if j + 1 < n and sql[j + 1] == c:
+                        j += 2
+                        continue
+                    break
+                j += 1
+            out.append(sql[i:j + 1])
+            i = j + 1
+        elif c == "$" and (m := re.match(r"\$([A-Za-z_][A-Za-z_0-9]*)?\$", sql[i:])):
+            tag = m.group(0)
+            j = sql.find(tag, i + len(tag))
+            j = n if j < 0 else j + len(tag)
+            out.append(sql[i:j])
+            i = j
+        elif sql.startswith("--", i):
+            j = sql.find("\n", i)
+            j = n if j < 0 else j
+            out.append(sql[i:j])
+            i = j
+        elif sql.startswith("/*", i):
+            j = sql.find("*/", i + 2)
+            j = n if j < 0 else j + 2
+            out.append(sql[i:j])
+            i = j
+        elif sql.startswith("%s", i):
+            k += 1
+            out.append(f"${k}")
+            i += 2
+        elif sql.startswith("%%", i):
+            out.append("%")
+            i += 2
+        else:
+            out.append(c)
+            i += 1
+    return "".join(out)
+
+
+class _Diag:
+    def __init__(self, e: BaseException):
+        self.message_primary = getattr(e, "message", None) or str(e)
+        self.message_detail = getattr(e, "detail", None)
+        self.message_hint = getattr(e, "hint", None)
+        self.sqlstate = getattr(e, "sqlstate", None)
+
+
+class AsyncpgError(psycopg.Error):
+    """An asyncpg PostgreSQL error, as the SDK reads psycopg's."""
+
+    def __init__(self, e: BaseException):
+        self.cause = e  # (first: psycopg's constructor reads sqlstate)
+        super().__init__(getattr(e, "message", None) or str(e))
+
+    @property
+    def sqlstate(self) -> Optional[str]:  # type: ignore[override]
+        return getattr(self.cause, "sqlstate", None)
+
+    @property
+    def diag(self) -> _Diag:  # type: ignore[override]
+        return _Diag(self.cause)
+
+
+class _Cursor:
+    def __init__(self, conn: Any, dicts: bool):
+        self.conn, self.dicts = conn, dicts
+        self.description: Optional[list[str]] = None
+        self.rowcount = -1
+        self.statusmessage: Optional[str] = None
+        self._rows: list = []
+
+    async def __aenter__(self) -> "_Cursor":
+        return self
+
+    async def __aexit__(self, *exc) -> bool:
+        return False
+
+    async def execute(self, sql: str, params: Any = None, prepare: Any = None) -> None:
+        q = to_dollar(sql) if params else sql
+        args = list(params or ())
+
+        async def run():
+            ps = await self.conn.prepare(q)
+            return ps, await ps.fetch(*args)
+
+        try:
+            if self.conn.is_in_transaction():
+                ps, rows = await run()
+            else:
+                async with self.conn.transaction():
+                    ps, rows = await run()
+        except Exception as e:  # asyncpg.PostgresError carries sqlstate
+            if getattr(e, "sqlstate", None):
+                raise AsyncpgError(e) from e
+            raise
+        attrs = list(ps.get_attributes())
+        self.description = [a.name for a in attrs] or None
+        self._rows = [dict(r) if self.dicts else tuple(r) for r in rows] if attrs else []
+        self.statusmessage = ps.get_statusmsg()
+        m = re.search(r"(\d+)\s*$", self.statusmessage or "")
+        self.rowcount = int(m.group(1)) if m else len(self._rows)
+
+    async def fetchall(self) -> list:
+        return self._rows
+
+    async def fetchone(self) -> Any:
+        return self._rows[0] if self._rows else None
+
+
+class _Conn:
+    def __init__(self, conn: Any):
+        self._c = conn
+
+    def transaction(self):
+        return self._c.transaction()
+
+    def cursor(self, row_factory: Any = None) -> _Cursor:
+        return _Cursor(self._c, row_factory is not None)
+
+
+class asyncpg_source:
+    """A login for AsyncGnomon from an asyncpg pool (a connection per call) or one asyncpg
+    connection (used one call at a time)."""
+
+    def __init__(self, source: Any):
+        self.source = source
+        self.lock = asyncio.Lock()
+
+    @contextlib.asynccontextmanager
+    async def connection(self) -> AsyncIterator[_Conn]:
+        s = self.source
+        if hasattr(s, "acquire"):
+            async with s.acquire() as c:
+                yield _Conn(c)
+        else:
+            async with self.lock:
+                yield _Conn(s)

gnomon/client.py

@@ -0,0 +1,361 @@
+"""The gnomon SDK, sync (the specification is sdk/spec/SPEC.md).
+
+Generated from ``aio.py`` by ``tools/unasync.py``: do not edit; change aio.py.
+
+    gn = Gnomon(app="postgresql://app@db/app", runs="postgresql://runs@db/app")
+    with gn.start_run("Refund order 3", actor="support-bot") as run:
+        run.sql("update orders set status = 'refunded' where id = 3")
+        summary = run.summary()
+        if summary["checks"]["passed"] and summary["changes"]["rows"] <= 20:
+            run.merge()
+        else:
+            run.submit("Refund order 3")
+"""
+from __future__ import annotations
+
+import threading
+import time
+import contextlib
+import json
+from typing import Any, Iterator, Callable, Optional
+
+import psycopg
+import psycopg.conninfo
+from psycopg.rows import dict_row
+
+from ._util import parse_json, split_sql
+from .errors import (GnomonError, IncompatibleVersionError, MergeConflictError, SerializationError,
+                     to_gnomon_error)
+
+SUPPORTED_API = (2, 2)
+
+
+class Adapter:
+    """A login: a connection string (a connection per call), a psycopg connection (used one call
+    at a time), or a pool with a ``connection()`` context manager (psycopg_pool)."""
+
+    def __init__(self, source: Any):
+        self.source = source
+        self.lock = threading.Lock()
+
+    @contextlib.contextmanager
+    def connection(self) -> Iterator[psycopg.Connection]:
+        s = self.source
+        if isinstance(s, str):
+            conn = psycopg.Connection.connect(s, autocommit=True)
+            try:
+                yield conn
+            finally:
+                conn.close()
+        elif hasattr(s, "connection") and not isinstance(s, psycopg.Connection):
+            with s.connection() as conn:
+                yield conn
+        else:
+            with self.lock:
+                yield s
+
+    def query(self, sql: str, params: tuple = ()) -> list[dict]:
+        with self.connection() as conn:
+            with conn.cursor(row_factory=dict_row) as cur:
+                cur.execute(sql, params)
+                return cur.fetchall() if cur.description else []
+
+
+class Gnomon:
+    """The application's side: starts runs, reads them, accepts them; review; core."""
+
+    def __init__(self, app: Any, runs: Any = None, *, statement_timeout_ms: int = 30000, lock_timeout_ms: int = 5000,
+                 max_rows: int = 1000, on_statement: Optional[Callable[[dict], None]] = None, skip_version_check: bool = False,
+                 url: Optional[str] = None):
+        # url: the database's connection string (host, port, database), for run.url() and
+        # run.connection(); default: the app login's, when it is a string
+        self.url = url or (app if isinstance(app, str) else None)
+        self.app = Adapter(app)
+        self.runs = Adapter(runs) if runs is not None else None
+        self.statement_timeout_ms, self.lock_timeout_ms, self.max_rows = statement_timeout_ms, lock_timeout_ms, max_rows
+        self.on_statement = on_statement
+        self._checked = skip_version_check
+        self.requests = Requests(self)
+        self.branches = Branches(self)
+
+    def ready(self) -> None:
+        """Check once that the extension speaks this SDK's API."""
+        if self._checked:
+            return
+        try:
+            rows = self.app.query("select gnomon.api_version() as v")
+        except psycopg.Error as e:
+            raise IncompatibleVersionError("this server's gnomon has no gnomon.api_version(): it predates this SDK (needs API 2)", cause=e)
+        v = rows[0]["v"]
+        if not SUPPORTED_API[0] <= v <= SUPPORTED_API[1]:
+            raise IncompatibleVersionError(f"gnomon API {v} on the server; this SDK speaks {SUPPORTED_API[0]}..{SUPPORTED_API[1]}")
+        self._checked = True
+
+    def _event(self, login: str, run: Optional[str], sql: str, ok: bool, error: Optional[GnomonError] = None) -> None:
+        if self.on_statement:
+            self.on_statement({"login": login, "run": run, "statement": sql, "ok": ok, "error": error})
+
+    def query(self, sql: str, params: tuple = ()) -> list[dict]:
+        """A statement on the app login, errors typed."""
+        self.ready()
+        try:
+            rows = self.app.query(sql, params)
+        except psycopg.Error as e:
+            ge = to_gnomon_error(e)
+            self._event("app", None, sql, False, ge)
+            raise ge from e
+        self._event("app", None, sql, True)
+        return rows
+
+    def value(self, sql: str, params: tuple = ()) -> Any:
+        rows = self.query(sql, params)
+        return next(iter(rows[0].values())) if rows else None
+
+    def json(self, sql: str, params: tuple = ()) -> Any:
+        return parse_json(self.value(f"select ({sql})::text", params))
+
+    def start_run(self, task: Optional[str] = None, *, actor: Optional[str] = None, kind: str = "logged",
+                        ttl: Optional[str] = None, parent: str = "main") -> "Run":
+        if self.runs is None:
+            raise GnomonError("start_run needs the runs login: Gnomon(app=..., runs=...)")
+        branch = self.value("select gnomon.run_start(%s, %s, %s, %s::interval, %s)", (task, actor, kind, ttl, parent))
+        return Run(self, branch, task=task, actor=actor)
+
+    def run(self, branch: str, *, task: Optional[str] = None, actor: Optional[str] = None) -> "Run":
+        return Run(self, branch, task=task, actor=actor)
+
+    # ---- core
+    def status(self) -> dict:
+        return self.json("select gnomon.status()")
+
+    def doctor(self) -> list[dict]:
+        return self.query("select check_name, status, detail from gnomon.doctor()")
+
+    def describe(self, schema: Optional[str] = None) -> dict:
+        return self.json("select gnomon.describe(%s)", (schema,))
+
+    def diff(self, a: str, b: Optional[str] = None) -> list[dict]:
+        rows = self.query("select tbl::text as table, key::text as key, kind, col as column, old::text as old, new::text as new from gnomon.diff(%s, %s)", (a, b))
+        return [{**r, "key": parse_json(r["key"]), "old": parse_json(r["old"]), "new": parse_json(r["new"])} for r in rows]
+
+    def history(self, table: str, key: Any, *, branch: str = "main") -> list[dict]:
+        rows = self.query("select vkey, xid::text as xid, role, op, at, actor, reason, row_image::text as row_image, on_branch from gnomon.history(%s, %s::jsonb, %s)",
+                                (table, json.dumps(key), branch))
+        return [{**r, "row_image": parse_json(r["row_image"])} for r in rows]
+
+    def blame(self, table: str, key: Any, *, branch: str = "main") -> list[dict]:
+        rows = self.query("select col, value::text as value, xid::text as xid, role, actor, reason, at, on_branch from gnomon.blame(%s, %s::jsonb, %s)",
+                                (table, json.dumps(key), branch))
+        return [{**r, "value": parse_json(r["value"])} for r in rows]
+
+    def revert(self, table: str, at: str, *, keys: Optional[list] = None, filter: Optional[str] = None,
+                     columns: Optional[list] = None, branch: str = "main", dry_run: bool = False) -> dict:
+        return self.json("select gnomon.revert(%s, %s, %s::jsonb, %s, %s::text[], %s, %s)",
+                               (table, at, json.dumps(keys) if keys is not None else None, filter, columns, branch, dry_run))
+
+
+class Run:
+    """One agent run: a branch, its envelope, and what the application does with it."""
+
+    def __init__(self, gnomon: Gnomon, branch: str, *, task: Optional[str] = None, actor: Optional[str] = None):
+        self.gnomon, self.branch, self.task, self.actor = gnomon, branch, task, actor
+        self.state = "open"
+        self._seq = 0
+
+    def __enter__(self) -> "Run":
+        return self
+
+    def __exit__(self, *exc) -> None:
+        if self.state == "open":
+            self.discard()
+
+    def _check_open(self) -> None:
+        if self.state != "open":
+            raise GnomonError(f"run {self.branch} is {self.state}", code="branch_closed")
+
+    def _envelope(self, cur, reason: Optional[str], timeout_ms: Optional[int]) -> None:
+        g = self.gnomon
+        self._seq += 1
+        cur.execute("select gnomon.begin(%s)", (self.branch,))
+        cur.execute("select gnomon.note(%s, %s, %s::jsonb)", (self.actor or "agent", reason,
+                                                                       json.dumps({"run": self.branch, "call": self._seq, "task": self.task})))
+        cur.execute("select set_config('statement_timeout', %s, true), set_config('lock_timeout', %s, true)",
+                          (str(timeout_ms or g.statement_timeout_ms), str(g.lock_timeout_ms)))
+
+    def sql(self, statement: str, params: tuple = (), *, reason: Optional[str] = None, timeout_ms: Optional[int] = None,
+                  max_rows: Optional[int] = None) -> dict:
+        """One statement on the run (the envelope). Rows past the cap: ``truncated``. Retried on a
+        serialization failure (the envelope held nothing else)."""
+        self._check_open()
+        g = self.gnomon
+        g.ready()
+        if len(split_sql(statement)) > 1:
+            raise GnomonError("one statement per call: this one holds several (split it, or use transaction())", code="one_statement", sqlstate="42601")
+        cap = max_rows or g.max_rows
+        for attempt in range(4):
+            try:
+                with g.runs.connection() as conn:
+                    with conn.transaction():
+                        with conn.cursor(row_factory=dict_row) as cur:
+                            self._envelope(cur, reason, timeout_ms)
+                            # (the extended protocol: psycopg sends parameters, so PostgreSQL refuses a second statement)
+                            cur.execute(statement, params or None, prepare=False)
+                            rows = cur.fetchall() if cur.description else []
+                            rowcount = cur.rowcount
+                g._event("runs", self.branch, statement, True)
+                return {"rows": rows[:cap], "row_count": rowcount, "truncated": len(rows) > cap}
+            except psycopg.Error as e:
+                ge = to_gnomon_error(e)
+                g._event("runs", self.branch, statement, False, ge)
+                if isinstance(ge, SerializationError) and attempt < 3:
+                    time.sleep(0.02 * 2 ** attempt)
+                    continue
+                raise ge from e
+
+    @contextlib.contextmanager
+    def transaction(self, *, reason: Optional[str] = None, timeout_ms: Optional[int] = None) -> Iterator[Any]:
+        """Several statements in one envelope: ``with run.transaction() as tx: tx.execute(...)``."""
+        self._check_open()
+        g = self.gnomon
+        g.ready()
+        try:
+            with g.runs.connection() as conn:
+                with conn.transaction():
+                    with conn.cursor(row_factory=dict_row) as cur:
+                        self._envelope(cur, reason, timeout_ms)
+                        yield cur
+        except psycopg.Error as e:
+            raise to_gnomon_error(e) from e
+
+    def diff(self) -> list[dict]:
+        return self.gnomon.diff(self.branch)
+
+    def summary(self, *, run_checks: bool = True) -> dict:
+        return self.gnomon.json("select gnomon.run_summary(%s, %s)", (self.branch, run_checks))
+
+    def checks(self) -> list[dict]:
+        return self.gnomon.query("select check_name, passed, detail from gnomon.run_checks(%s)", (self.branch,))
+
+    def merge(self, *, close: bool = True) -> dict:
+        g = self.gnomon
+        mid = g.value("select gnomon.merge(%s, null, 'atomic', null, null, null, %s)", (self.branch, close))
+        m = g.json("select to_jsonb(m) from gnomon.merge m where id = %s", (mid,))
+        if m["status"] == "conflicts":
+            c = g.json("select coalesce(jsonb_agg(to_jsonb(c)), '[]') from gnomon.conflict c where merge_id = %s", (mid,))
+            raise MergeConflictError(f"merging run {self.branch} stopped on {len(c)} conflict(s)", mid, c)
+        self.state = "merged"
+        return {"merge": mid, "status": m["status"], "stats": m["stats"]}
+
+    def submit(self, title: Optional[str] = None, body: Optional[str] = None) -> "Request":
+        cr = self.gnomon.value("select gnomon.open_change_request(%s, %s, %s)",
+                                     (self.branch, title or self.task or f"Run {self.branch}", body))
+        self.state = "submitted"
+        return Request(self.gnomon, cr)
+
+    def discard(self) -> None:
+        if self.state == "discarded":
+            return
+        self.gnomon.query("select gnomon.run_end(%s, false)", (self.branch,))
+        self.state = "discarded"
+
+    def url(self, base: Optional[str] = None) -> str:
+        """A connection string pinned to the run, for a tool that wants one: a login of the run's own
+        (gnomon.run_login), confined to the run's branch, dropped when the run ends."""
+        self._check_open()
+        b = base or self.gnomon.url
+        if not b:
+            raise GnomonError("run.url() needs a base connection string (the url option, or an argument)", code="run_url_base")
+        login = self.gnomon.json("select gnomon.run_login(%s)", (self.branch,))
+        return psycopg.conninfo.make_conninfo(b, user=login["role"], password=login["password"])
+
+    def connection(self, base: Optional[str] = None) -> psycopg.Connection:
+        """A connection pinned to the run (close it when done)."""
+        return psycopg.Connection.connect(self.url(base))
+
+    def checkpoint(self, name: str) -> str:
+        return self.gnomon.value("select gnomon.checkpoint(%s, %s)", (f"{self.branch}:{name}", self.branch))
+
+    def rewind(self, name: str) -> dict:
+        return self.gnomon.json("select gnomon.rewind(%s, %s)", (self.branch, f"{self.branch}:{name}"))
+
+    def dry_run(self, statements: list[str], *, checks: bool = True) -> dict:
+        return self.gnomon.json("select gnomon.dry_run(%s::text[], %s, %s)", (statements, self.branch, checks))
+
+    def describe(self, schema: Optional[str] = None) -> dict:
+        with self.gnomon.runs.connection() as conn:
+            with conn.transaction():
+                with conn.cursor() as cur:
+                    cur.execute("select gnomon.begin(%s)", (self.branch,))
+                    cur.execute("select gnomon.describe(%s)::text", (schema,))
+                    return parse_json((cur.fetchone())[0])
+
+
+class Requests:
+    def __init__(self, g: Gnomon):
+        self.gnomon = g
+
+    def open(self) -> list["Request"]:
+        rows = self.gnomon.query("select id from gnomon.change_request where status in ('open', 'approved') order by id")
+        return [Request(self.gnomon, r["id"]) for r in rows]
+
+    def get(self, id: int) -> "Request":
+        return Request(self.gnomon, id)
+
+
+class Request:
+    """A change request: what a product's own Approve, Reject and Undo buttons call."""
+
+    def __init__(self, g: Gnomon, id: int):
+        self.gnomon, self.id = g, id
+
+    def status(self) -> dict:
+        return self.gnomon.json("select gnomon.change_request_status(%s)", (self.id,))
+
+    def diff(self) -> list[dict]:
+        return self.gnomon.diff((self.status())["branch"])
+
+    def approve(self, body: Optional[str] = None) -> None:
+        self.gnomon.query("select gnomon.review(%s, 'approve', %s)", (self.id, body))
+
+    def reject(self, note: Optional[str] = None) -> None:
+        self.gnomon.query("select gnomon.review(%s, 'reject', %s)", (self.id, note))
+
+    def comment(self, body: str, *, table: Optional[str] = None, key: Any = None, column: Optional[str] = None) -> int:
+        return self.gnomon.value("select gnomon.comment(%s, %s, %s::regclass, %s::jsonb, %s)",
+                                       (self.id, body, table, json.dumps(key) if key is not None else None, column))
+
+    def merge(self, *, close: bool = True) -> int:
+        return self.gnomon.value("select gnomon.merge(%s, null, 'atomic', null, null, null, %s)", ((self.status())["branch"], close))
+
+    def undo(self) -> dict:
+        s = self.status()
+        if not s.get("merge"):
+            raise GnomonError(f"change request {self.id} is not merged")
+        return self.gnomon.json("select gnomon.undo_merge(%s)", (s["merge"]["id"],))
+
+    def wait(self, timeout_s: float = 300) -> dict:
+        return self.gnomon.json("select gnomon.wait_for_review(%s, %s::interval)", (self.id, f"{timeout_s} seconds"))
+
+
+class Branches:
+    def __init__(self, g: Gnomon):
+        self.gnomon = g
+
+    def create(self, name: str, *, parent: str = "main", kind: str = "logged", at: Optional[str] = None) -> str:
+        self.gnomon.query("select gnomon.create_branch(%s, %s, %s, %s)", (name, parent, kind, at))
+        return name
+
+    def scratch(self, prefix: str = "test", parent: str = "main") -> str:
+        return self.gnomon.value("select gnomon.create_scratch(%s, %s)", (prefix, parent))
+
+    def drop(self, name: str) -> None:
+        self.gnomon.query("select gnomon.drop_branch(%s, true)", (name,))
+
+    def rebase(self, name: str, *, dry_run: bool = False) -> dict:
+        return self.gnomon.json("select gnomon.rebase(%s, null, %s)", (name, dry_run))
+
+    def list(self) -> list[dict]:
+        return self.gnomon.query("select * from gnomon.branches order by created_at")
+
+    def options(self, name: str, *, c_module: bool = False) -> str:
+        return self.gnomon.value("select gnomon.connection_options(%s, %s)", (name, c_module))

gnomon/django.py

@@ -0,0 +1,68 @@
+"""Django integration.
+
+Add ``"gnomon.django.BranchMiddleware"`` to MIDDLEWARE. A request with an ``X-Gnomon-Branch``
+header (or a ``request.gnomon_branch`` attribute set by an earlier middleware) runs on that branch.
+
+* Default: the request's connection is put on the branch with gnomon.use() before the view and
+  reset to main afterwards, so a persistent connection never carries the branch to the next request.
+* ``GNOMON_TRANSACTION_POOLING = True`` (behind pgbouncer in transaction mode): every transaction
+  of the request starts with gnomon.begin(), which lasts exactly one transaction. This needs
+  ``ATOMIC_REQUESTS = True`` so the request's queries share a transaction.
+
+``GNOMON_BRANCH_HEADER`` changes the header name.
+"""
+from __future__ import annotations
+
+import contextvars
+
+from django.conf import settings
+from django.core.exceptions import ImproperlyConfigured
+from django.db import connections
+
+_current = contextvars.ContextVar("gnomon_branch", default=None)
+
+
+def current_branch():
+    """The branch of the request being served, or None."""
+    return _current.get()
+
+
+class _BeginEachTransaction:
+    def __init__(self, conn, branch):
+        self.conn, self.branch, self.done = conn, branch, set()
+
+    def __call__(self, execute, sql, params, many, context):
+        if not self.conn.in_atomic_block:
+            raise ImproperlyConfigured("GNOMON_TRANSACTION_POOLING needs ATOMIC_REQUESTS = True")
+        key = id(self.conn.atomic_blocks[0])
+        if key not in self.done:
+            self.done.add(key)
+            context["cursor"].execute("select gnomon.begin(%s)", [self.branch])
+        return execute(sql, params, many, context)
+
+
+class BranchMiddleware:
+    def __init__(self, get_response):
+        self.get_response = get_response
+        name = getattr(settings, "GNOMON_BRANCH_HEADER", "X-Gnomon-Branch")
+        self.header = "HTTP_" + name.upper().replace("-", "_")
+
+    def __call__(self, request):
+        b = getattr(request, "gnomon_branch", None) or request.META.get(self.header)
+        if not b or b == "main":
+            return self.get_response(request)
+        conn = connections["default"]
+        token = _current.set(b)
+        try:
+            if getattr(settings, "GNOMON_TRANSACTION_POOLING", False):
+                with conn.execute_wrapper(_BeginEachTransaction(conn, b)):
+                    return self.get_response(request)
+            with conn.cursor() as cur:
+                cur.execute("select gnomon.use(%s)", [b])
+            try:
+                return self.get_response(request)
+            finally:
+                with conn.cursor() as cur:
+                    cur.execute("select gnomon.use('main')")
+        finally:
+            _current.reset(token)

gnomon/errors.py

@@ -0,0 +1,65 @@
+"""Typed errors (sdk/spec/SPEC.md, Errors). The SQLSTATE stays the standard one; the slug in the
+hint (``[gnomon:<slug>]``) is gnomon's stable code."""
+from __future__ import annotations
+
+import re
+from typing import Any, Optional
+
+_SLUG = re.compile(r"\[gnomon:([a-z_]+)\]")
+
+
+class GnomonError(Exception):
+    def __init__(self, message: str, *, code: Optional[str] = None, sqlstate: Optional[str] = None,
+                 hint: Optional[str] = None, cause: Optional[BaseException] = None):
+        super().__init__(message)
+        self.code, self.sqlstate, self.hint, self.cause = code, sqlstate, hint, cause
+
+
+class SerializationError(GnomonError): ...
+class PermissionError(GnomonError): ...  # noqa: A001 (gnomon's, not the builtin)
+class BranchClosedError(GnomonError): ...
+class StatementTimeoutError(GnomonError): ...
+class UnsupportedOnBranchError(GnomonError): ...
+class NotRoutedError(GnomonError): ...
+class OverQuotaError(GnomonError): ...
+class IncompatibleVersionError(GnomonError): ...
+
+
+class ChecksFailedError(GnomonError):
+    results: Any = None
+
+
+class MergeConflictError(GnomonError):
+    def __init__(self, message: str, merge_id: int, conflicts: list):
+        super().__init__(message, code="merge_conflict")
+        self.merge_id, self.conflicts = merge_id, conflicts
+
+
+def to_gnomon_error(err: BaseException) -> GnomonError:
+    """Map a psycopg error to a typed gnomon error."""
+    if isinstance(err, GnomonError):
+        return err
+    sqlstate = getattr(err, "sqlstate", None)
+    diag = getattr(err, "diag", None)
+    hint = getattr(diag, "message_hint", None) if diag is not None else None
+    msg = getattr(diag, "message_primary", None) or str(err)
+    m = _SLUG.search(f"{hint or ''} {msg}")
+    slug = m.group(1) if m else None
+    kw = dict(code=slug, sqlstate=sqlstate, hint=hint, cause=err)
+    if sqlstate == "40001":
+        return SerializationError(msg, **kw)
+    if sqlstate == "57014":
+        return StatementTimeoutError(msg, **kw)
+    if sqlstate == "42P10" or slug in ("sql_only_dml", "upsert_on_branch"):
+        return UnsupportedOnBranchError(msg + " (on a branch, write an upsert as UPDATE then INSERT ... WHERE NOT EXISTS, or load gnomon's C module)", **kw)
+    if slug == "not_routed":
+        return NotRoutedError(msg, **kw)
+    if sqlstate == "25006" or slug == "branch_closed":
+        return BranchClosedError(msg, **kw)
+    if sqlstate == "53400" or slug == "over_quota":
+        return OverQuotaError(msg, **kw)
+    if sqlstate == "23514":
+        return ChecksFailedError(msg, **kw)
+    if sqlstate == "42501":
+        return PermissionError(msg, **kw)
+    return GnomonError(msg, **kw)