pgcommon 0.1.0__py3-none-any.whl

pgcommon/guc.py

@@ -0,0 +1,142 @@
+"""RLS tenant identity (`GUCSet`) and its request-scoped propagation via `contextvars`.
+
+Go carries the active identity as an explicit `context.Context` value threaded through every
+call. Python has no equivalent explicit parameter — `contextvars.ContextVar` is the mechanism
+used instead, and it is scoped per `asyncio` task (each request gets its own copy once a task
+is spawned for it). See ARCHITECTURE.md for the exact set/reset timing this requires in ASGI
+middleware — getting that wrong is the single most likely place for a cross-request GUC leak.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Awaitable, Iterator
+from contextlib import contextmanager
+from contextvars import ContextVar, Token
+from dataclasses import dataclass, field
+from typing import Protocol
+
+from pgcommon.errors import EmptyTenantRoleError, InconsistentGUCSetError, InvalidTenantRoleError
+
+__all__ = [
+    "GUCSet",
+    "guc_set_from_context",
+    "inject_gucs",
+    "with_guc_set",
+    "with_validated_guc_set",
+]
+
+
+@dataclass(frozen=True, slots=True)
+class GUCSet:
+    """The RLS identity injected as `app.user_id` / `app.tenant_id` / `app.tenant_roles`."""
+
+    user_id: str = ""
+    tenant_id: str = ""
+    tenant_roles: tuple[str, ...] = field(default_factory=tuple)
+
+    def is_zero(self) -> bool:
+        """True when every field is empty — the pool skips GUC injection entirely in this case."""
+        return self.user_id == "" and self.tenant_id == "" and not self.tenant_roles
+
+    def validate(self) -> None:
+        """Raise if this GUCSet is internally inconsistent.
+
+        Checks, in order: `user_id` set without `tenant_id` (InconsistentGUCSetError), a role
+        containing a comma — the GUC serialization separator (InvalidTenantRoleError), and an
+        empty-string role (EmptyTenantRoleError).
+        """
+        if self.user_id and not self.tenant_id:
+            raise InconsistentGUCSetError("GUCSet.user_id is set but tenant_id is empty")
+        for role in self.tenant_roles:
+            if "," in role:
+                raise InvalidTenantRoleError(f"GUCSet.tenant_roles entry {role!r} contains a comma")
+            if role == "":
+                raise EmptyTenantRoleError("GUCSet.tenant_roles contains an empty-string entry")
+
+
+_current_guc_set: ContextVar[GUCSet | None] = ContextVar("pgcommon_guc_set", default=None)
+
+
+@contextmanager
+def with_guc_set(guc_set: GUCSet) -> Iterator[None]:
+    """Store `guc_set` on the current context for the duration of the `with` block.
+
+    Does not validate — use `with_validated_guc_set` in auth middleware so an inconsistent
+    identity is rejected before it reaches connection acquisition.
+    """
+    token: Token[GUCSet | None] = _current_guc_set.set(guc_set)
+    try:
+        yield
+    finally:
+        _current_guc_set.reset(token)
+
+
+@contextmanager
+def with_validated_guc_set(guc_set: GUCSet) -> Iterator[None]:
+    """Like `with_guc_set`, but calls `guc_set.validate()` first.
+
+    Mirrors the Go README's `WithValidatedGUCSet` — intended for auth middleware, which should
+    catch the raised error and respond 401 rather than let an inconsistent identity reach the
+    database layer.
+    """
+    guc_set.validate()
+    with with_guc_set(guc_set):
+        yield
+
+
+def guc_set_from_context() -> GUCSet | None:
+    """Return the `GUCSet` stored by the innermost enclosing `with_guc_set`, or `None`.
+
+    Pass this directly as `Config.guc_provider` for automatic RLS injection.
+    """
+    return _current_guc_set.get()
+
+
+class _GUCExecutable(Protocol):
+    """Structural shape shared by `asyncpg.Connection` and `asyncpg.pool.PoolConnectionProxy`."""
+
+    def execute(self, query: str, *args: object) -> Awaitable[str]: ...
+
+
+_INJECT_SQL = (
+    "SELECT set_config('app.user_id', $1, {is_local}), "
+    "set_config('app.tenant_id', $2, {is_local}), "
+    "set_config('app.tenant_roles', $3, {is_local})"
+)
+
+
+async def inject_gucs(conn: _GUCExecutable, guc_set: GUCSet) -> None:
+    """Session-scoped GUC injection (`is_local=false`) — always session-scoped, regardless of
+    `Config.pgbouncer_mode`. This is for manual, one-off injection; the pool's automatic
+    injection respects the configured mode (see `pgcommon.tx` for the PgBouncer-mode path).
+
+    A zero `GUCSet` (see `GUCSet.is_zero`) is silently skipped — no GUCs are set, and RLS
+    policies fall through to their own default (typically deny).
+    """
+    guc_set.validate()
+    if guc_set.is_zero():
+        return
+    await conn.execute(
+        _INJECT_SQL.format(is_local="false"),
+        guc_set.user_id,
+        guc_set.tenant_id,
+        ",".join(guc_set.tenant_roles),
+    )
+
+
+async def _inject_gucs_local(conn: _GUCExecutable, guc_set: GUCSet) -> None:
+    """Transaction-local GUC injection (`is_local=true`) for PgBouncer transaction-pooling mode.
+
+    Must be called inside an already-open transaction — `is_local=true` settings are cleared
+    automatically on `COMMIT` / `ROLLBACK`, which is exactly what makes this safe when the
+    backend connection is handed back to PgBouncer afterward.
+    """
+    guc_set.validate()
+    if guc_set.is_zero():
+        return
+    await conn.execute(
+        _INJECT_SQL.format(is_local="true"),
+        guc_set.user_id,
+        guc_set.tenant_id,
+        ",".join(guc_set.tenant_roles),
+    )

pgcommon/integrations/__init__.py

@@ -0,0 +1,5 @@
+"""Optional framework glue. Nothing under `pgcommon.integrations` is imported by `pgcommon`
+itself — each submodule pulls in its own framework only when explicitly imported, so the core
+library stays framework-agnostic. `pgcommon.integrations.fastapi` requires the `pgcommon[fastapi]`
+extra.
+"""

pgcommon/integrations/fastapi.py

@@ -0,0 +1,151 @@
+"""FastAPI integration: pool lifespan wiring, GUC-injection middleware, and a `/healthz` example.
+
+Optional — requires the `pgcommon[fastapi]` extra. The core library (`pgcommon.*`) never
+imports `fastapi` or `starlette`; this is the only module in the package that does, and
+`pgcommon/__init__.py` does not import it — importing `pgcommon.integrations.fastapi` is
+always an explicit, opt-in step for a consuming service.
+
+## contextvars and request isolation — read this before wiring `GUCAuthMiddleware`
+
+pgcommon propagates the active `GUCSet` via `contextvars.ContextVar`, not an explicit context
+parameter threaded through every call the way Go's `context.Context` is. Under ASGI, each
+inbound request is handled by its own `asyncio.Task`, and task creation captures a *copy* of
+the current `contextvars.Context` at the moment the task starts — so a `ContextVar.set()` call
+made while handling one request is invisible to a concurrently-running request's task, even
+though both run on the same event loop. That per-task copy is what makes it safe to stash a
+`GUCSet` in middleware without a lock or an explicit per-request object.
+
+The failure mode to avoid: **never call `ContextVar.set()` without capturing and later
+`reset()`-ing its token in the same middleware invocation.** `GUCAuthMiddleware` below uses
+`with_validated_guc_set(...)` — a context manager — wrapped tightly around `call_next(request)`
+for exactly this reason: the token is reset in a `finally` block the moment the response for
+*this* request has been produced, so the GUCSet cannot leak into code that runs later on the
+same task (a connection-pool cleanup callback, a background job scheduled from the handler, or
+— on a pathological ASGI server that reuses tasks across requests — the next request entirely).
+Setting a `ContextVar` and forgetting to reset it is the single most likely place for a
+cross-request GUC leak in this stack, more so than in Go, where the value is an explicit
+argument and cannot silently outlive its caller.
+"""
+
+from __future__ import annotations
+
+from collections.abc import AsyncIterator, Awaitable, Callable
+from contextlib import AbstractAsyncContextManager, asynccontextmanager
+from typing import Protocol
+
+from fastapi import Depends, FastAPI, Request
+from fastapi.responses import JSONResponse
+from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint
+from starlette.responses import Response
+from starlette.types import ASGIApp
+
+from pgcommon.config import Config
+from pgcommon.errors import EmptyTenantRoleError, InconsistentGUCSetError, InvalidTenantRoleError
+from pgcommon.guc import GUCSet, with_validated_guc_set
+from pgcommon.pool import Pool
+
+__all__ = [
+    "GUCAuthMiddleware",
+    "IdentityExtractor",
+    "create_lifespan",
+    "get_pool",
+    "healthz",
+]
+
+
+class IdentityExtractor(Protocol):
+    """Extracts the caller's identity from a request. Implement this for your auth scheme
+    (JWT claims, session lookup, an upstream `platform-gincommon`-style request context, ...).
+
+    Return `None` for requests with no identity (unauthenticated, internal tooling) — no GUCSet
+    is stored and downstream RLS policies apply their own default (typically deny-all).
+    """
+
+    def __call__(self, request: Request) -> GUCSet | None: ...
+
+
+def create_lifespan(
+    config: Config,
+) -> Callable[[FastAPI], AbstractAsyncContextManager[None]]:
+    """Build an ASGI lifespan that opens the `Pool` once at startup and drains it at shutdown.
+
+    The pool is stored on `app.state.pool` — never reconstructed per request. Usage::
+
+        app = FastAPI(lifespan=create_lifespan(config))
+    """
+
+    @asynccontextmanager
+    async def lifespan(app: FastAPI) -> AsyncIterator[None]:
+        pool = await Pool.create(config)
+        app.state.pool = pool
+        try:
+            yield
+        finally:
+            await pool.drain_and_close()
+
+    return lifespan
+
+
+def get_pool(request: Request) -> Pool:
+    """FastAPI dependency yielding the shared `Pool` built by `create_lifespan`.
+
+    Usage in a route handler::
+
+        @app.get("/orders")
+        async def list_orders(pool: Pool = Depends(get_pool)):
+            return await pool.with_conn(...)
+    """
+    return request.app.state.pool  # type: ignore[no-any-return]
+
+
+class GUCAuthMiddleware(BaseHTTPMiddleware):
+    """Extracts a `GUCSet` via `identity_extractor` and scopes it to this request only.
+
+    Responds `401` if the extracted identity fails `GUCSet.validate()` (e.g. a user ID without
+    a tenant ID, or a role containing a comma) — rejecting an inconsistent identity here, before
+    it ever reaches connection acquisition, is exactly what the Go README's
+    `WithValidatedGUCSet` does in the equivalent middleware example.
+    """
+
+    def __init__(self, app: ASGIApp, identity_extractor: IdentityExtractor) -> None:
+        super().__init__(app)
+        self._identity_extractor = identity_extractor
+
+    async def dispatch(self, request: Request, call_next: RequestResponseEndpoint) -> Response:
+        guc_set = self._identity_extractor(request)
+        if guc_set is None:
+            return await call_next(request)
+
+        try:
+            with with_validated_guc_set(guc_set):
+                return await call_next(request)
+        except (InconsistentGUCSetError, InvalidTenantRoleError, EmptyTenantRoleError):
+            return JSONResponse({"detail": "invalid identity"}, status_code=401)
+
+
+def healthz(*, utilization_threshold: float = 0.9) -> Callable[..., Awaitable[JSONResponse]]:
+    """Build a `/healthz` route handler with utilization-based load shedding.
+
+    Returns `503` when the database is unreachable, or when the pool's utilization is at or
+    above `utilization_threshold` (default 90%) — shedding new requests before the pool is
+    fully exhausted protects the connections already serving in-flight requests. Usage::
+
+        app.add_api_route("/healthz", healthz(), methods=["GET"])
+    """
+
+    async def handler(pool: Pool = Depends(get_pool)) -> JSONResponse:
+        status = await pool.health()
+        if not status.healthy:
+            return JSONResponse({"detail": "database unavailable"}, status_code=503)
+        if status.utilization >= utilization_threshold:
+            return JSONResponse({"detail": "service overloaded"}, status_code=503)
+        return JSONResponse(
+            {
+                "healthy": status.healthy,
+                "acquired_conns": status.acquired_conns,
+                "max_conns": status.max_conns,
+                "utilization": round(status.utilization, 3),
+            }
+        )
+
+    return handler

pgcommon/logger.py

@@ -0,0 +1,71 @@
+"""The `Logger` protocol pgcommon logs through, plus two reference adapters.
+
+Mirrors the Go repo's `port.Logger` interface. `pgcommon` never imports a concrete logging
+framework at the top level — callers wire in whichever adapter fits, or write their own by
+implementing the four methods below.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Any, Protocol, runtime_checkable
+
+__all__ = ["Logger", "StdlibLogger", "StructlogLogger"]
+
+
+@runtime_checkable
+class Logger(Protocol):
+    """Structured logger contract used throughout pgcommon (slow-query warnings, migration steps)."""
+
+    def debug(self, msg: str, /, **fields: object) -> None: ...
+
+    def info(self, msg: str, /, **fields: object) -> None: ...
+
+    def warn(self, msg: str, /, **fields: object) -> None: ...
+
+    def error(self, msg: str, /, **fields: object) -> None: ...
+
+
+class StdlibLogger:
+    """Adapts `logging.Logger` to the `Logger` protocol; fields are passed via the `extra` dict."""
+
+    def __init__(self, logger: logging.Logger | None = None) -> None:
+        self._logger = logger if logger is not None else logging.getLogger("pgcommon")
+
+    def debug(self, msg: str, /, **fields: object) -> None:
+        self._logger.debug(msg, extra={"pgcommon_fields": fields} if fields else None)
+
+    def info(self, msg: str, /, **fields: object) -> None:
+        self._logger.info(msg, extra={"pgcommon_fields": fields} if fields else None)
+
+    def warn(self, msg: str, /, **fields: object) -> None:
+        self._logger.warning(msg, extra={"pgcommon_fields": fields} if fields else None)
+
+    def error(self, msg: str, /, **fields: object) -> None:
+        self._logger.error(msg, extra={"pgcommon_fields": fields} if fields else None)
+
+
+class StructlogLogger:
+    """Adapts a `structlog` logger to the `Logger` protocol — the reference implementation
+    mirroring the Go repo's `NewZapLogger`. Requires the `pgcommon[structlog]` extra."""
+
+    def __init__(self, logger: Any | None = None) -> None:
+        try:
+            import structlog
+        except ImportError as exc:  # pragma: no cover - exercised only without the extra
+            raise ImportError(
+                "StructlogLogger requires the 'structlog' package — install pgcommon[structlog]"
+            ) from exc
+        self._logger = logger if logger is not None else structlog.get_logger("pgcommon")
+
+    def debug(self, msg: str, /, **fields: object) -> None:
+        self._logger.debug(msg, **fields)
+
+    def info(self, msg: str, /, **fields: object) -> None:
+        self._logger.info(msg, **fields)
+
+    def warn(self, msg: str, /, **fields: object) -> None:
+        self._logger.warning(msg, **fields)
+
+    def error(self, msg: str, /, **fields: object) -> None:
+        self._logger.error(msg, **fields)

pgcommon/metrics.py

@@ -0,0 +1,179 @@
+"""Prometheus metrics registration and the OpenTelemetry `db.query` span helper.
+
+Mirrors the Go repo's `pkg/pgmetrics`. Both Prometheus and OTel are optional: the pool,
+transactions, and GUC injection all work without ever calling `init` or supplying
+`Config.tracer`.
+"""
+
+from __future__ import annotations
+
+import time
+from dataclasses import dataclass
+
+from opentelemetry.trace import SpanKind, Status, StatusCode, Tracer
+from prometheus_client import REGISTRY as _DEFAULT_REGISTRY
+from prometheus_client import CollectorRegistry, Counter, Gauge, Histogram
+
+from pgcommon._internal.query_event import QueryEvent
+from pgcommon._internal.sql_verb import sql_operation
+
+__all__ = ["Metrics", "OTelQueryRecorder", "current", "init", "init_with_registry"]
+
+
+@dataclass(frozen=True, slots=True)
+class Metrics:
+    """The handles `init` / `init_with_registry` return — pass to `OTelQueryRecorder` or use directly."""
+
+    query_total: Counter
+    query_duration: Histogram
+    pool_acquire_total: Counter
+    pool_acquire_duration: Histogram
+    slow_query_total: Counter
+    retry_total: Counter
+    retry_exhausted_total: Counter
+    build_info: Gauge
+
+
+_metrics: Metrics | None = None
+
+
+def current() -> Metrics | None:
+    """Return the `Metrics` handle registered by `init` / `init_with_registry`, or `None`.
+
+    Used internally by `Pool`, `run_in_tx`, and the retry helpers to record metrics only when
+    observability has been opted into — none of them require it to function.
+    """
+    return _metrics
+
+
+def init(service: str, version: str) -> Metrics:
+    """Register Prometheus metrics against the default registry. Idempotent — first caller wins."""
+    return init_with_registry(service, version, _DEFAULT_REGISTRY)
+
+
+def init_with_registry(service: str, version: str, registry: CollectorRegistry) -> Metrics:
+    """Register Prometheus metrics against a specific registry (use a fresh one in tests).
+
+    Idempotent per process: subsequent calls (against any registry) return the handles from
+    the first call rather than re-registering, mirroring the Go repo's `sync.Once` guard.
+    """
+    global _metrics
+    if _metrics is not None:
+        return _metrics
+
+    build_info = Gauge(
+        "pgcommon_build_info",
+        "Static build metadata for the service using pgcommon; value is always 1.",
+        ["service", "version"],
+        registry=registry,
+    )
+    build_info.labels(service=service, version=version).set(1)
+
+    _metrics = Metrics(
+        query_total=Counter(
+            "pgcommon_query_total",
+            "Queries executed, labeled by SQL verb and outcome.",
+            ["operation", "status"],
+            registry=registry,
+        ),
+        query_duration=Histogram(
+            "pgcommon_query_duration_seconds",
+            "Query duration in seconds, labeled by SQL verb.",
+            ["operation"],
+            registry=registry,
+        ),
+        pool_acquire_total=Counter(
+            "pgcommon_pool_acquire_total",
+            "Connection acquire attempts, labeled by outcome.",
+            ["status"],
+            registry=registry,
+        ),
+        pool_acquire_duration=Histogram(
+            "pgcommon_pool_acquire_duration_seconds",
+            "Time spent waiting to acquire a connection from the pool.",
+            registry=registry,
+        ),
+        slow_query_total=Counter(
+            "pgcommon_slow_query_total",
+            "Queries exceeding Config.slow_query_threshold, counted regardless of log rate-limiting.",
+            registry=registry,
+        ),
+        retry_total=Counter(
+            "pgcommon_retry_total",
+            "Retried transaction attempts, labeled by retry reason.",
+            ["reason"],
+            registry=registry,
+        ),
+        retry_exhausted_total=Counter(
+            "pgcommon_retry_exhausted_total",
+            "Retry sequences that exhausted all attempts without success, labeled by reason.",
+            ["reason"],
+            registry=registry,
+        ),
+        build_info=build_info,
+    )
+    return _metrics
+
+
+class OTelQueryRecorder:
+    """Records one `db.query` span per query, reconstructed after the fact from `elapsed`.
+
+    asyncpg's query logger fires once the query has already completed, so there is no live
+    span to start before the query and end after it (pgx's `QueryTracer` gets both a start and
+    an end hook). Instead this back-dates the span's start time using `record.elapsed`, which
+    yields the correct duration and still nests correctly under the caller's active span
+    because `asyncio.Task`/`call_soon` scheduling preserves the `contextvars` context that
+    carries OTel's current-span pointer.
+    """
+
+    def __init__(
+        self,
+        tracer: Tracer,
+        *,
+        allow_full_statements: bool = False,
+        metrics: Metrics | None = None,
+    ) -> None:
+        self._tracer = tracer
+        self._allow_full_statements = allow_full_statements
+        self._metrics = metrics
+
+    def allow_full_sql(self) -> None:
+        """Opt in to including the full SQL text in `db.statement` (default: verb only)."""
+        self._allow_full_statements = True
+
+    def __call__(self, record: QueryEvent) -> None:
+        """Register this as an asyncpg query logger: `conn.add_query_logger(recorder)`."""
+        end = time.time()
+        start = end - record.elapsed
+        operation = sql_operation(record.query)
+
+        attributes: dict[str, str | int] = {
+            "db.system": "postgresql",
+            "db.statement": record.query if self._allow_full_statements else operation,
+        }
+        params = record.conn_params
+        user = getattr(params, "user", None)
+        host = getattr(params, "host", None)
+        port = getattr(params, "port", None)
+        if user:
+            attributes["db.user"] = user
+        if host:
+            attributes["net.peer.name"] = host
+        if port:
+            attributes["net.peer.port"] = port
+
+        span = self._tracer.start_span(
+            "db.query",
+            kind=SpanKind.CLIENT,
+            attributes=attributes,
+            start_time=int(start * 1_000_000_000),
+        )
+        status = "error" if record.exception is not None else "success"
+        if record.exception is not None:
+            span.record_exception(record.exception)
+            span.set_status(Status(StatusCode.ERROR))
+        span.end(end_time=int(end * 1_000_000_000))
+
+        if self._metrics is not None:
+            self._metrics.query_total.labels(operation=operation, status=status).inc()
+            self._metrics.query_duration.labels(operation=operation).observe(record.elapsed)

pgcommon/migrate/__init__.py

@@ -0,0 +1,8 @@
+"""Migration runner — a minimal, dependency-light golang-migrate equivalent over plain .sql files."""
+
+from __future__ import annotations
+
+from pgcommon.migrate.errors import MigrationDirtyError, NoMigrationsFoundError
+from pgcommon.migrate.runner import Migration, Runner
+
+__all__ = ["Migration", "MigrationDirtyError", "NoMigrationsFoundError", "Runner"]

pgcommon/migrate/__main__.py

@@ -0,0 +1,6 @@
+import sys
+
+from pgcommon.migrate.cli import main
+
+if __name__ == "__main__":
+    sys.exit(main())

pgcommon/migrate/cli.py

@@ -0,0 +1,78 @@
+"""`pgcommon-migrate` console script — also runnable as `python -m pgcommon.migrate`.
+
+Intended for a Kubernetes init container or a CI/CD pre-deploy job — never as a sidecar
+running alongside the live application, and never invoked from every replica's own boot path
+simultaneously. See the README's migration deployment-pattern table.
+"""
+
+from __future__ import annotations
+
+import argparse
+import asyncio
+import os
+import sys
+
+from pgcommon.errors import MigrationDirtyError, NoMigrationsFoundError
+from pgcommon.logger import StdlibLogger
+from pgcommon.migrate.runner import Runner
+
+__all__ = ["main"]
+
+
+def _build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(
+        prog="pgcommon-migrate", description="Apply or roll back pgcommon SQL migrations."
+    )
+    parser.add_argument("--dsn", default=None, help="Postgres DSN (default: $DATABASE_URL)")
+    parser.add_argument(
+        "--source",
+        default="migrations",
+        help="Directory of *.up.sql / *.down.sql files (default: ./migrations)",
+    )
+
+    subparsers = parser.add_subparsers(dest="command", required=True)
+    subparsers.add_parser("up", help="Apply all pending migrations")
+
+    down_parser = subparsers.add_parser("down", help="Roll back N steps (default: 1)")
+    down_parser.add_argument("n", type=int, nargs="?", default=1)
+
+    steps_parser = subparsers.add_parser("steps", help="Apply +N or roll back -N steps")
+    steps_parser.add_argument("n", type=int)
+
+    subparsers.add_parser("version", help="Print the current schema version")
+    return parser
+
+
+async def _run(args: argparse.Namespace) -> int:
+    dsn = args.dsn or os.environ.get("DATABASE_URL")
+    if not dsn:
+        print("error: --dsn or DATABASE_URL is required", file=sys.stderr)
+        return 2
+
+    runner = Runner(dsn, args.source, logger=StdlibLogger())
+    try:
+        if args.command == "up":
+            await runner.up()
+        elif args.command == "down":
+            await runner.down(args.n)
+        elif args.command == "steps":
+            await runner.steps(args.n)
+        elif args.command == "version":
+            version, dirty = await runner.version()
+            print(f"version={version} dirty={dirty}")
+    except MigrationDirtyError as exc:
+        print(f"error: {exc}", file=sys.stderr)
+        return 1
+    except NoMigrationsFoundError as exc:
+        print(f"error: {exc}", file=sys.stderr)
+        return 1
+    return 0
+
+
+def main(argv: list[str] | None = None) -> int:
+    args = _build_parser().parse_args(argv)
+    return asyncio.run(_run(args))
+
+
+if __name__ == "__main__":  # pragma: no cover
+    sys.exit(main())

pgcommon/migrate/errors.py

@@ -0,0 +1,12 @@
+"""Re-exports of the migration-specific sentinel errors defined in `pgcommon.errors`.
+
+Kept as a separate import path (`pgcommon.migrate.errors`) to mirror the Go repo's package
+layout; the classes themselves live in `pgcommon.errors` so `isinstance` checks against
+errors raised anywhere else in the library still work.
+"""
+
+from __future__ import annotations
+
+from pgcommon.errors import MigrationDirtyError, NoMigrationsFoundError
+
+__all__ = ["MigrationDirtyError", "NoMigrationsFoundError"]