walbox 1.0.0__py3-none-any.whl

walbox/__init__.py

@@ -0,0 +1,27 @@
+"""walbox: a PostgreSQL logical replication client."""
+
+from walbox.abc import ChangeEvent
+from walbox.abc import ChangeKind
+from walbox.abc import CheckpointHandle
+from walbox.abc import Metrics
+from walbox.abc import MetricsCallback
+from walbox.abc import Transaction
+from walbox.abc import WalboxOptions
+from walbox.builder import Walbox
+from walbox.checkpoint import ConnectionPool
+from walbox.client import Client
+from walbox.errors import WalboxError
+
+__all__ = [
+    "ChangeEvent",
+    "ChangeKind",
+    "CheckpointHandle",
+    "Client",
+    "ConnectionPool",
+    "Metrics",
+    "MetricsCallback",
+    "Transaction",
+    "Walbox",
+    "WalboxError",
+    "WalboxOptions",
+]

walbox/abc.py

@@ -0,0 +1,184 @@
+"""Shared data types and protocols for walbox."""
+
+import time
+from collections.abc import Callable
+from dataclasses import dataclass
+from dataclasses import field
+from enum import StrEnum
+from typing import Any
+from typing import Protocol
+
+from psycopg import AsyncConnection
+
+from walbox.errors import CheckpointError
+
+
+class ChangeKind(StrEnum):
+    """The kind of row-level change a `ChangeEvent` carries."""
+
+    INSERT = "insert"
+    UPDATE = "update"
+    DELETE = "delete"
+    TRUNCATE = "truncate"
+
+
+@dataclass
+class ChangeEvent:
+    """A single row-level change within a transaction.
+
+    `new` and `old` depend on `kind`:
+
+    - INSERT: `new` has the inserted row, `old` is None.
+    - UPDATE: `new` has the row after the change, `old` has it before.
+      `old` needs a usable REPLICA IDENTITY (a primary key, or
+      REPLICA IDENTITY FULL) on the table; without one it's None.
+    - DELETE: `new` is None, `old` has the deleted row (same REPLICA
+      IDENTITY requirement as UPDATE).
+    - TRUNCATE: both are None.
+
+    Dict keys are column names; values are the column's Python-typed
+    value, with SQL NULL represented as None.
+    """
+
+    kind: ChangeKind
+    table: str
+    new: dict[str, Any] | None = None
+    old: dict[str, Any] | None = None
+
+
+class CheckpointStore(Protocol):
+    """Durably tracks the last replayed LSN so replication can resume after restart."""
+
+    async def load(self) -> int | None:
+        """Return the last durably saved LSN, or None if none has been saved yet."""
+        ...
+
+    async def save(
+        self,
+        lsn: int,
+        *,
+        connection: AsyncConnection[Any] | None = None,
+    ) -> None:
+        """Durably persist `lsn` as the new replay position.
+
+        `connection`, if given, is an already-open Postgres connection or
+        transaction the implementation can join instead of opening its own.
+        """
+        ...
+
+
+@dataclass(frozen=True, slots=True)
+class CheckpointHandle:
+    """A `CheckpointStore` bound to one run, handed to every handler call."""
+
+    _store: CheckpointStore
+    _on_saved: Callable[[int, float], None] | None = None
+    _max_lsn: int | None = None
+
+    async def save(
+        self,
+        lsn: int,
+        *,
+        connection: AsyncConnection[Any] | None = None,
+    ) -> None:
+        """Durably persist `lsn` via the bound `CheckpointStore`.
+
+        `connection`, if given, is an already-open Postgres connection or
+        transaction the underlying store can join instead of opening its own.
+
+        Raises:
+            CheckpointError: If `lsn` is greater than the commit LSN of the
+                transaction this handle was constructed for. Saving a
+                checkpoint ahead of what was actually processed risks
+                PostgreSQL recycling WAL for data walbox never handled.
+        """
+        if self._max_lsn is not None and lsn > self._max_lsn:
+            message = (
+                f"refusing to save checkpoint lsn={lsn}: greater than "
+                f"the dispatched transaction's commit_lsn={self._max_lsn}"
+            )
+            raise CheckpointError(message)
+        started_at = time.monotonic()
+        await self._store.save(lsn, connection=connection)
+        if self._on_saved is not None:
+            self._on_saved(lsn, time.monotonic() - started_at)
+
+
+@dataclass(frozen=True, slots=True)
+class Transaction:
+    """A committed transaction and the row-level changes it contains."""
+
+    xid: int
+    commit_lsn: int
+    commit_time: int
+    changes: list[ChangeEvent] = field(default_factory=list)
+
+
+@dataclass(frozen=True)
+class Metrics:
+    """A point-in-time snapshot of replication counters and gauges.
+
+    Handed to `WalboxOptions.on_metrics` on the same timer as status
+    updates. No historical aggregation (rolling windows, percentiles,
+    rates) is done here; that's on the application if it wants one.
+    """
+
+    consumer_name: str
+    receive_lsn: int
+    checkpoint_lsn: int
+    replication_lag_bytes: int
+    transactions_processed: int
+    changes_processed: int
+    reconnect_count: int
+    last_handler_latency_seconds: float
+    queue_depth: int
+    last_keepalive_at: float
+    last_checkpoint_latency_seconds: float
+    transactions_since_checkpoint: int
+
+
+MetricsCallback = Callable[[Metrics], None]
+
+
+@dataclass
+class WalboxOptions:
+    """Configuration for a walbox client, shared by every checkpoint backend.
+
+    `Client` takes this plus a `CheckpointStore` directly;
+    `Walbox` constructs the checkpoint store for you.
+
+    `on_metrics` runs synchronously on the same timer as status updates.
+    Don't block in it; hand the `Metrics` snapshot off to a queue or task
+    if you need to send it elsewhere. See docs/production/monitoring.md.
+    """
+
+    consumer_name: str
+    dsn: str
+    slot_name: str
+    publication_name: str
+
+    max_pending_transactions: int = 100
+    status_interval: int = 10
+    on_metrics: MetricsCallback | None = None
+
+    def __post_init__(self) -> None:
+        """Validate required fields, raising `ValueError` if any are invalid.
+
+        Raises:
+            ValueError: If any required field is missing or invalid.
+        """
+        for name, value in (
+            ("consumer_name", self.consumer_name),
+            ("dsn", self.dsn),
+            ("slot_name", self.slot_name),
+            ("publication_name", self.publication_name),
+        ):
+            if not value or not value.strip():
+                message = f"{name} must not be blank"
+                raise ValueError(message)
+        if self.max_pending_transactions <= 0:
+            message = "max_pending_transactions must be > 0"
+            raise ValueError(message)
+        if self.status_interval <= 0:
+            message = "status_interval must be > 0"
+            raise ValueError(message)

walbox/builder.py

@@ -0,0 +1,50 @@
+"""Builds a fully wired `Client` from public `WalboxOptions`."""
+
+from walbox.abc import WalboxOptions
+from walbox.checkpoint import ConnectionPool
+from walbox.checkpoint import PostgresCheckpointStore
+from walbox.client import Client
+
+
+class Walbox:
+    """Constructs a `Client` without exposing checkpoint-store wiring."""
+
+    @staticmethod
+    def build(options: WalboxOptions) -> Client:
+        """Build a client whose checkpoint store opens ad hoc connections.
+
+        Opens and closes a new connection for every `checkpoint.save()`
+        call. Prefer `build_with_pool` unless you have a reason not to add
+        the `psycopg-pool` dependency; this is here for that case, and for
+        checkpoint volume low enough that the per-call connection cost
+        doesn't matter.
+
+        Returns:
+            A `Client` ready to `run()`.
+        """
+        checkpoint_store = PostgresCheckpointStore(
+            options.dsn,
+            consumer_name=options.consumer_name,
+        )
+        return Client(options, checkpoint_store)
+
+    @staticmethod
+    def build_with_pool(
+        options: WalboxOptions,
+        pool: ConnectionPool,
+    ) -> Client:
+        """Build a client whose checkpoint store reuses your own connection pool.
+
+        `pool` is owned by the caller: opening and closing it (for example
+        via `async with AsyncConnectionPool(...) as pool:`) is your
+        responsibility, not walbox's. The same pool can also be reused for a
+        handler's own downstream writes.
+
+        Returns:
+            A `Client` ready to `run()`.
+        """
+        checkpoint_store = PostgresCheckpointStore.from_pool(
+            pool,
+            consumer_name=options.consumer_name,
+        )
+        return Client(options, checkpoint_store)

walbox/checkpoint.py

@@ -0,0 +1,191 @@
+"""CheckpointStore implementations for durably tracking replay position."""
+
+import contextlib
+import logging
+import time
+from collections.abc import AsyncGenerator
+from collections.abc import Callable
+from contextlib import AbstractAsyncContextManager
+from typing import Any
+from typing import Protocol
+
+import psycopg
+from psycopg import AsyncConnection
+from psycopg import sql
+
+from walbox.errors import CheckpointError
+
+logger = logging.getLogger("walbox.checkpoint")
+
+_Acquire = Callable[[], AbstractAsyncContextManager[AsyncConnection[Any]]]
+
+
+class ConnectionPool(Protocol):
+    """Structural shape of a Postgres connection pool.
+
+    Matches `psycopg_pool.AsyncConnectionPool` (an async context manager
+    that checks out a connection and returns it on exit), but is never
+    imported from `psycopg_pool`. Any object shaped like this works, so
+    `PostgresCheckpointStore.from_pool` adds no dependency beyond `psycopg`.
+    """
+
+    def connection(self) -> AbstractAsyncContextManager[AsyncConnection[Any]]:
+        """Check out a connection, returning it to the pool on block exit."""
+        ...
+
+
+@contextlib.asynccontextmanager
+async def _connect(dsn: str) -> AsyncGenerator[AsyncConnection[Any]]:
+    async with await psycopg.AsyncConnection.connect(dsn) as conn:
+        yield conn
+
+
+class PostgresCheckpointStore:
+    """A `CheckpointStore` backed by a row in a Postgres table.
+
+    `save(connection=...)` can join a caller-supplied connection's
+    transaction instead of opening its own, letting an application commit
+    its own sink write and the checkpoint update atomically in one
+    transaction.
+
+    Without `connection=`, `load()` and `save()` open one ad hoc connection
+    per call, which is fine for checkpointing's low call volume. Use
+    `from_pool` instead to reuse a connection pool the application already
+    manages.
+    """
+
+    def __init__(
+        self,
+        dsn: str,
+        *,
+        consumer_name: str,
+        table: str = "walbox_checkpoint",
+    ) -> None:
+        """Initialize with the DSN to connect with and the consumer to track.
+
+        `table` must be a trusted, developer-supplied identifier, never
+        end-user input: Postgres doesn't allow parameterizing table names,
+        so it's composed into SQL via `sql.Identifier` instead of a bind
+        parameter.
+        """
+        self._acquire: _Acquire = lambda: _connect(dsn)
+        self._configure(consumer_name=consumer_name, table=table)
+
+    @classmethod
+    def from_pool(
+        cls,
+        pool: ConnectionPool,
+        *,
+        consumer_name: str,
+        table: str = "walbox_checkpoint",
+    ) -> "PostgresCheckpointStore":
+        """Build a store whose ad hoc `load()`/`save()` calls reuse `pool`.
+
+        `pool` is owned by the application; it's never connected to or
+        closed here. This only changes where connections for `load()` and
+        connection-less `save()` calls come from. `save(lsn, connection=...)`
+        already uses whatever connection the caller passes in, pool or not.
+
+        Returns:
+            A `PostgresCheckpointStore` backed by `pool`.
+        """
+        store = cls.__new__(cls)
+        store._acquire = pool.connection  # ruff: ignore[private-member-access]: alternate constructor
+        store._configure(consumer_name=consumer_name, table=table)  # ruff: ignore[private-member-access]
+        return store
+
+    def _configure(self, *, consumer_name: str, table: str) -> None:
+        self._consumer_name = consumer_name
+        self._table = sql.Identifier(table)
+        self._schema_ready = False
+
+    async def load(self) -> int | None:
+        """Return the last durably saved LSN, or None if this consumer has none yet.
+
+        Raises:
+            CheckpointError: If the saved LSN is negative.
+        """
+        started_at = time.monotonic()
+        async with self._acquire() as conn:
+            await self._ensure_schema(conn)
+            await conn.commit()
+            query = sql.SQL(
+                "SELECT lsn FROM {table} WHERE consumer_name = %s",
+            ).format(table=self._table)
+            cursor = await conn.execute(query, (self._consumer_name,))
+            row = await cursor.fetchone()
+            lsn = row[0] if row is not None else None
+        if lsn is not None and lsn < 0:
+            message = (
+                f"checkpoint for consumer {self._consumer_name!r} is negative: {lsn}"
+            )
+            raise CheckpointError(message)
+        logger.debug(
+            "checkpoint load completed in %.6fs, lsn=%s",
+            time.monotonic() - started_at,
+            lsn,
+            extra={"lsn": lsn},
+        )
+        return lsn
+
+    async def save(
+        self,
+        lsn: int,
+        *,
+        connection: AsyncConnection[Any] | None = None,
+    ) -> None:
+        """Durably persist `lsn` as the new replay position for this consumer.
+
+        If `connection` is given, the upsert (and, on first use, the
+        backing table's creation) runs on it and is left uncommitted, so the
+        caller's own commit makes it durable atomically with whatever else
+        it writes on that connection. Without `connection`, this acquires
+        its own connection and commits immediately.
+        """
+        started_at = time.monotonic()
+        if connection is not None:
+            await self._ensure_schema(connection)
+            await self._upsert(connection, lsn)
+        else:
+            async with self._acquire() as conn:
+                await self._ensure_schema(conn)
+                await self._upsert(conn, lsn)
+                await conn.commit()
+        logger.debug(
+            "checkpoint save completed in %.6fs, lsn=%s",
+            time.monotonic() - started_at,
+            lsn,
+            extra={"lsn": lsn},
+        )
+
+    async def _ensure_schema(self, conn: AsyncConnection[Any]) -> None:
+        """Create the backing table if needed, without committing.
+
+        `CREATE TABLE IF NOT EXISTS` is transactional in Postgres, so this
+        is safe to run on a caller-supplied `save(connection=...)` connection
+        and leave for the caller's own commit. Committing here would commit
+        the caller's in-progress transaction early, breaking the
+        same-transaction guarantee `connection=` exists to provide.
+        """
+        if self._schema_ready:
+            return
+        await conn.execute(
+            sql.SQL(
+                "CREATE TABLE IF NOT EXISTS {table} ("
+                "consumer_name TEXT PRIMARY KEY, "
+                "lsn BIGINT NOT NULL, "
+                "updated_at TIMESTAMPTZ NOT NULL DEFAULT now()"
+                ")",
+            ).format(table=self._table),
+        )
+        self._schema_ready = True
+
+    async def _upsert(self, conn: AsyncConnection[Any], lsn: int) -> None:
+        await conn.execute(
+            sql.SQL(
+                "INSERT INTO {table} (consumer_name, lsn) VALUES (%s, %s) "
+                "ON CONFLICT (consumer_name) DO UPDATE SET lsn = EXCLUDED.lsn, "
+                "updated_at = now()",
+            ).format(table=self._table),
+            (self._consumer_name, lsn),
+        )