feather-cdc-postgres 0.1.0__py3-none-any.whl

feather_cdc_postgres/ownership.py

@@ -0,0 +1,103 @@
+"""Source-side advisory ownership for one registration (#3079 D8).
+
+Session-scoped `pg_try_advisory_lock` on the SOURCE database, keyed on the slot: Postgres
+releases it when the owning backend dies, so a killed drainer cannot split-brain. Held across
+EVERY protected phase (publication reconcile, snapshot, backfill, each drain batch) — call
+`assert_held()` before each; a dead connection or a lost lock raises `OwnershipLost`.
+"""
+from __future__ import annotations
+
+import psycopg
+
+# NOT renamed with the module. This string is hashed into the advisory-lock key on the SOURCE
+# database: two drainers of one slot exclude each other only while they agree on it. A deployment
+# mid-cutover runs the old tree and this one against the same source, so changing it here would
+# let both hold "the" lock at once. It changes when every drainer changes, not when the import
+# name does.
+KEY_PREFIX = "pg_md_cdc:"
+
+# `pg_try_advisory_lock(bigint)` splits its key across pg_locks: classid = the high 32 bits,
+# objid = the low 32 bits, objsubid = 1 (2 marks the two-int form). Probed on PG 17,
+# 04-Sep-2026 — a negative hashtext sign-extends, so classid is 4294967295, never 0. Counting
+# any granted advisory lock instead would call someone else's lock our own.
+_HELD_SQL = """
+select count(*) from pg_locks
+where locktype = 'advisory' and pid = pg_backend_pid() and granted
+  and objsubid = 1
+  and classid = ((hashtext(%s)::bigint >> 32) & 4294967295)::bigint::oid
+  and objid = (hashtext(%s)::bigint & 4294967295)::bigint::oid
+"""
+
+
+class OwnershipLost(RuntimeError):
+    """The owner session is gone (connection dead or advisory lock no longer held)."""
+
+
+class OwnerSession:
+    def __init__(self, dsn: str, slot: str, *, connect_timeout: int = 10) -> None:
+        self.dsn = dsn
+        self.slot = slot
+        self.key = KEY_PREFIX + slot
+        self._conn: psycopg.Connection | None = None
+        self._held = False
+
+    # --- lifecycle ---
+    def acquire(self) -> bool:
+        if self._held:
+            return True
+        self._conn = psycopg.connect(self.dsn, connect_timeout=10, autocommit=True)
+        got = self._conn.execute(
+            "select pg_try_advisory_lock(hashtext(%s)::bigint)", (self.key,)
+        ).fetchone()[0]
+        self._held = bool(got)
+        if not self._held:
+            self._close()
+        return self._held
+
+    def assert_held(self) -> None:
+        """Cheap liveness + ownership re-check. Doubles as the keepalive that stops a managed
+        Postgres (Railway) reaping the owner connection as idle during a long destination load."""
+        if not self._held or self._conn is None or self._conn.closed:
+            raise OwnershipLost(f"owner session for slot {self.slot!r} is not held")
+        try:
+            self._conn.execute("select 1").fetchone()
+            granted = self._conn.execute(_HELD_SQL, (self.key, self.key)).fetchone()[0]
+        except Exception as exc:  # noqa: BLE001 — a dead connection IS a lost lock
+            self._held = False
+            raise OwnershipLost(
+                f"owner session for slot {self.slot!r} died: {type(exc).__name__}: {exc}"
+            ) from exc
+        if not granted:
+            self._held = False
+            raise OwnershipLost(
+                f"advisory lock for slot {self.slot!r} is no longer held by this session"
+            )
+
+    def release(self) -> None:
+        if self._conn is not None and not self._conn.closed and self._held:
+            try:
+                self._conn.execute("select pg_advisory_unlock(hashtext(%s)::bigint)", (self.key,))
+            except Exception:  # noqa: BLE001 — closing the session releases it anyway
+                pass
+        self._held = False
+        self._close()
+
+    def _close(self) -> None:
+        if self._conn is not None:
+            try:
+                self._conn.close()
+            except Exception:  # noqa: BLE001
+                pass
+            self._conn = None
+
+    # --- context manager ---
+    def __enter__(self) -> "OwnerSession":
+        self.acquired = self.acquire()
+        return self
+
+    def __exit__(self, *exc) -> None:
+        self.release()
+
+    @property
+    def held(self) -> bool:
+        return self._held

feather_cdc_postgres/publication.py

@@ -0,0 +1,157 @@
+"""Publication + replica-identity reconciliation, always schema-qualified (#3079 D4/D6).
+
+The publication is an EXPLICIT table list reconciled to the desired set every run — never
+`FOR TABLES IN SCHEMA`, because a table that auto-joins gets no baseline from us. Relation
+identity is `(schema, table)` throughout: two schemas may legally hold the same table name.
+Replica-identity ALTERs are diff-based (they take a brief exclusive lock) and plannable.
+The publish list is `insert, update, delete, truncate`: a mirror that silently kept every row
+of a truncated source table would be stale bronze with no event to say so (#3079 review). An
+existing publication with a narrower list (the pilot's) is ALTERed up, never inherited.
+"""
+from __future__ import annotations
+
+from typing import Iterable, Sequence
+
+import psycopg
+from psycopg import sql
+
+Relation = tuple[str, str]
+
+
+def _connect(dsn: str) -> psycopg.Connection:
+    return psycopg.connect(dsn, connect_timeout=10, autocommit=True)
+
+
+def publication_exists(dsn: str, publication: str) -> bool:
+    with _connect(dsn) as conn:
+        return conn.execute(
+            "select 1 from pg_publication where pubname = %s", (publication,)
+        ).fetchone() is not None
+
+
+def published_relations(dsn: str, publication: str) -> set[Relation]:
+    with _connect(dsn) as conn:
+        return {
+            (r[0], r[1])
+            for r in conn.execute(
+                "select schemaname, tablename from pg_publication_tables where pubname = %s",
+                (publication,),
+            ).fetchall()
+        }
+
+
+def drop_publication(dsn: str, publication: str) -> None:
+    with _connect(dsn) as conn:
+        conn.execute(sql.SQL("drop publication if exists {}").format(sql.Identifier(publication)))
+
+
+PUBLISH = "insert, update, delete, truncate"
+
+_PUB_FLAGS = {"insert": "pubinsert", "update": "pubupdate", "delete": "pubdelete",
+              "truncate": "pubtruncate"}
+
+
+def reconcile(
+    dsn: str,
+    publication: str,
+    desired: Iterable[Relation],
+    *,
+    publish: str = PUBLISH,
+    plan_only: bool = False,
+) -> tuple[list[Relation], list[Relation]]:
+    """Bring `publication` to exactly `desired` (membership AND publish list). Returns
+    (added, dropped), both sorted."""
+    desired_set = {(s, t) for s, t in desired}
+    wanted_ops = {op.strip() for op in publish.split(",") if op.strip()}
+    with _connect(dsn) as conn:
+        flags = conn.execute(
+            "select pubinsert, pubupdate, pubdelete, pubtruncate from pg_publication "
+            "where pubname = %s", (publication,)
+        ).fetchone()
+        exists = flags is not None
+        # PG 18 refuses UPDATEs on a published REPLICA IDENTITY FULL table whose stored
+        # generated columns are not published, so the publication always publishes them.
+        gencols = conn.info.server_version >= 180000
+        if exists and not plan_only:
+            current_ops = {op for op, col in _PUB_FLAGS.items()
+                           if flags[list(_PUB_FLAGS).index(op)]}
+            if current_ops != wanted_ops:
+                conn.execute(sql.SQL("alter publication {} set (publish = {})").format(
+                    sql.Identifier(publication), sql.Literal(publish)))
+            if gencols and conn.execute(
+                "select pubgencols from pg_publication where pubname = %s", (publication,)
+            ).fetchone()[0] != "s":
+                conn.execute(sql.SQL("alter publication {} set (publish_generated_columns = stored)")
+                             .format(sql.Identifier(publication)))
+        current: set[Relation] = set()
+        if exists:
+            current = {
+                (r[0], r[1])
+                for r in conn.execute(
+                    "select schemaname, tablename from pg_publication_tables where pubname = %s",
+                    (publication,),
+                ).fetchall()
+            }
+        added = sorted(desired_set - current)
+        dropped = sorted(current - desired_set)
+        if plan_only:
+            return added, dropped
+        if not exists:
+            conn.execute(
+                sql.SQL("create publication {} with (publish = {}{})").format(
+                    sql.Identifier(publication), sql.Literal(publish),
+                    sql.SQL(", publish_generated_columns = stored" if gencols else ""),
+                )
+            )
+        for schema, table in added:
+            conn.execute(
+                sql.SQL("alter publication {} add table {}.{}").format(
+                    sql.Identifier(publication), sql.Identifier(schema), sql.Identifier(table)
+                )
+            )
+        for schema, table in dropped:
+            conn.execute(
+                sql.SQL("alter publication {} drop table {}.{}").format(
+                    sql.Identifier(publication), sql.Identifier(schema), sql.Identifier(table)
+                )
+            )
+    return added, dropped
+
+
+def ensure_replica_identity(
+    dsn: str,
+    relations: Sequence[Relation],
+    mode: str = "full",
+    *,
+    plan_only: bool = False,
+) -> list[Relation]:
+    """Diff-based: only tables whose `relreplident` differs from the declared mode are ALTERed.
+    `full` is what makes unchanged-TOAST updates decodable (see resolve_unchanged_toast)."""
+    if mode not in ("full", "default"):
+        raise ValueError(f"replica_identity must be 'full' or 'default', not {mode!r}")
+    want = "f" if mode == "full" else "d"
+    rels = list(relations)
+    if not rels:
+        return []
+    with _connect(dsn) as conn:
+        wanted = {(s_, t_) for s_, t_ in rels}
+        needs = [
+            (r[0], r[1])
+            for r in conn.execute(
+                "select n.nspname, c.relname from pg_class c "
+                "join pg_namespace n on n.oid = c.relnamespace "
+                "where c.relkind = 'r' and n.nspname = any(%s) and c.relreplident <> %s",
+                (sorted({s_ for s_, _ in rels}), want),
+            ).fetchall()
+            if (r[0], r[1]) in wanted
+        ]
+        if plan_only:
+            return sorted(needs)
+        for schema, table in needs:
+            conn.execute(
+                sql.SQL("alter table {}.{} replica identity {}").format(
+                    sql.Identifier(schema), sql.Identifier(table),
+                    sql.SQL("full" if mode == "full" else "default"),
+                )
+            )
+    return sorted(needs)