pgcommon 0.1.0__py3-none-any.whl

pgcommon/migrate/runner.py

@@ -0,0 +1,261 @@
+"""A minimal, dependency-light migration runner over plain `.sql` files.
+
+Deliberately not built on Alembic — Alembic assumes SQLAlchemy models, which this library does
+not have. This is a from-scratch, `golang-migrate`-equivalent runner: it reads
+`NNNN_description.up.sql` / `NNNN_description.down.sql` pairs from a directory, tracks the
+applied version in a `schema_migrations` table, and takes a Postgres advisory lock for the
+duration of the run so two runners started concurrently (e.g. two replicas booting at once)
+don't race each other.
+
+⚠️ **Never run migrations from every replica's boot path simultaneously.** The advisory lock
+here only serializes the runs — it does not make it safe for the application to start against
+a half-migrated schema, nor does it make concurrent runners fast or free of readiness-probe
+timeouts. Use an init container, a pre-deploy job, or a leader-only-on-boot pattern, and prefer
+additive (expand/contract) migrations so an old app version tolerates the new schema during a
+rolling deploy. See the README's migration deployment-pattern table.
+"""
+
+from __future__ import annotations
+
+import re
+import zlib
+from collections.abc import Awaitable, Callable
+from dataclasses import dataclass
+from pathlib import Path
+from typing import TypeVar
+
+import asyncpg
+
+from pgcommon.errors import MigrationDirtyError, NoMigrationsFoundError
+from pgcommon.logger import Logger
+
+__all__ = ["Migration", "Runner"]
+
+T = TypeVar("T")
+
+_FILENAME_RE = re.compile(r"^(?P<version>\d+)_(?P<description>.+)\.(?P<direction>up|down)\.sql$")
+_ADVISORY_LOCK_KEY = zlib.crc32(b"pgcommon-migrate")
+_LOCK_TIMEOUT_SECONDS = 30
+
+
+@dataclass(frozen=True, slots=True)
+class Migration:
+    """One migration version, with up and/or down SQL loaded from disk."""
+
+    version: int
+    description: str
+    up_sql: str | None
+    down_sql: str | None
+
+
+def _load_migrations(source: Path) -> dict[int, Migration]:
+    migrations: dict[int, Migration] = {}
+    try:
+        entries = sorted(source.iterdir())
+    except OSError as exc:
+        raise NoMigrationsFoundError(f"migration source {source} is not readable: {exc}") from exc
+
+    for path in entries:
+        if not path.is_file():
+            continue
+        m = _FILENAME_RE.match(path.name)
+        if not m:
+            continue
+        version = int(m.group("version"))
+        description = m.group("description")
+        direction = m.group("direction")
+        sql = path.read_text()
+
+        existing = migrations.get(version)
+        up_sql = existing.up_sql if existing else None
+        down_sql = existing.down_sql if existing else None
+        if direction == "up":
+            up_sql = sql
+        else:
+            down_sql = sql
+        migrations[version] = Migration(
+            version=version,
+            description=existing.description if existing else description,
+            up_sql=up_sql,
+            down_sql=down_sql,
+        )
+
+    if not migrations:
+        raise NoMigrationsFoundError(f"no *.up.sql / *.down.sql files found under {source}")
+    return migrations
+
+
+class Runner:
+    """Wraps a directory of SQL migrations and a target database.
+
+    `source` defaults are the caller's responsibility — pass an absolute or relative path to a
+    directory containing `NNNN_description.up.sql` / `NNNN_description.down.sql` pairs.
+    """
+
+    def __init__(self, dsn: str, source: str | Path, *, logger: Logger | None = None) -> None:
+        self._dsn = dsn
+        self._source = Path(source)
+        self._logger = logger
+
+    async def up(self) -> None:
+        """Apply all pending migrations, in ascending version order."""
+        await self._run(direction="up", steps=None)
+
+    async def down(self, n: int) -> None:
+        """Roll back the `n` most recently applied migrations, in descending version order."""
+        await self._run(direction="down", steps=n)
+
+    async def steps(self, n: int) -> None:
+        """Apply `n` migrations forward if `n >= 0`, or roll back `-n` if `n < 0`."""
+        if n >= 0:
+            await self._run(direction="up", steps=n)
+        else:
+            await self._run(direction="down", steps=-n)
+
+    async def version(self) -> tuple[int, bool]:
+        """Return `(version, dirty)` for the schema's current migration state."""
+        conn = await self._connect()
+        try:
+            return await self._with_advisory_lock(conn, lambda: self._version_locked(conn))
+        finally:
+            await conn.close()
+
+    async def _version_locked(self, conn: asyncpg.Connection) -> tuple[int, bool]:
+        await self._ensure_schema_migrations_table(conn)
+        return await self._current_version(conn)
+
+    async def _connect(self) -> asyncpg.Connection:
+        conn = await asyncpg.connect(self._dsn)
+        await conn.execute(f"SET lock_timeout = '{_LOCK_TIMEOUT_SECONDS}s'")
+        return conn
+
+    async def _ensure_schema_migrations_table(self, conn: asyncpg.Connection) -> None:
+        await conn.execute(
+            "CREATE TABLE IF NOT EXISTS schema_migrations ("
+            "version bigint PRIMARY KEY, dirty boolean NOT NULL DEFAULT false)"
+        )
+
+    async def _current_version(self, conn: asyncpg.Connection) -> tuple[int, bool]:
+        row = await conn.fetchrow("SELECT version, dirty FROM schema_migrations LIMIT 1")
+        if row is None:
+            return 0, False
+        return row["version"], row["dirty"]
+
+    async def _set_version(self, conn: asyncpg.Connection, version: int, *, dirty: bool) -> None:
+        async with conn.transaction():
+            await conn.execute("DELETE FROM schema_migrations")
+            await conn.execute(
+                "INSERT INTO schema_migrations (version, dirty) VALUES ($1, $2)", version, dirty
+            )
+
+    async def _with_advisory_lock(
+        self, conn: asyncpg.Connection, fn: Callable[[], Awaitable[T]]
+    ) -> T:
+        await conn.execute("SELECT pg_advisory_lock($1)", _ADVISORY_LOCK_KEY)
+        try:
+            return await fn()
+        finally:
+            await conn.execute("SELECT pg_advisory_unlock($1)", _ADVISORY_LOCK_KEY)
+
+    async def _apply_sql(
+        self, conn: asyncpg.Connection, sql: str, *, failing_version: int, success_version: int
+    ) -> None:
+        try:
+            async with conn.transaction():
+                await conn.execute(sql)
+        except Exception:
+            await self._set_version(conn, failing_version, dirty=True)
+            raise
+        await self._set_version(conn, success_version, dirty=False)
+
+    async def _run(self, *, direction: str, steps: int | None) -> None:
+        migrations = _load_migrations(self._source)
+        conn = await self._connect()
+        try:
+            await self._with_advisory_lock(
+                conn, lambda: self._run_locked(conn, migrations, direction, steps)
+            )
+        finally:
+            await conn.close()
+
+    async def _run_locked(
+        self,
+        conn: asyncpg.Connection,
+        migrations: dict[int, Migration],
+        direction: str,
+        steps: int | None,
+    ) -> None:
+        # CREATE TABLE IF NOT EXISTS is not safe against true concurrent creation in Postgres —
+        # two sessions racing this can both pass the "not exists" check and collide on the
+        # underlying pg_type entry. Only ever run it while holding the advisory lock.
+        await self._ensure_schema_migrations_table(conn)
+        current_version, dirty = await self._current_version(conn)
+        if dirty:
+            raise MigrationDirtyError(
+                f"schema is dirty at version {current_version} — a previous migration failed "
+                "mid-run; fix the migration and manually clear the dirty flag before retrying"
+            )
+
+        if direction == "up":
+            await self._run_up(conn, migrations, current_version, steps)
+        else:
+            assert steps is not None, "down/steps(negative) callers always pass an int"
+            await self._run_down(conn, migrations, current_version, steps)
+
+    async def _run_up(
+        self,
+        conn: asyncpg.Connection,
+        migrations: dict[int, Migration],
+        current_version: int,
+        steps: int | None,
+    ) -> None:
+        pending = sorted(v for v in migrations if v > current_version)
+        if steps is not None:
+            pending = pending[:steps]
+
+        if not pending:
+            if self._logger is not None:
+                self._logger.info("no pending migrations", version=current_version)
+            return
+
+        for version in pending:
+            migration = migrations[version]
+            if migration.up_sql is None:
+                raise NoMigrationsFoundError(f"missing *.up.sql for version {version}")
+            await self._apply_sql(
+                conn, migration.up_sql, failing_version=version, success_version=version
+            )
+            if self._logger is not None:
+                self._logger.info("applied migration", version=version, direction="up")
+
+    async def _run_down(
+        self,
+        conn: asyncpg.Connection,
+        migrations: dict[int, Migration],
+        current_version: int,
+        steps: int,
+    ) -> None:
+        applied = sorted(v for v in migrations if v <= current_version)
+        to_roll_back = applied[-steps:] if steps > 0 else []
+        if not to_roll_back:
+            if self._logger is not None:
+                self._logger.info("no migrations to roll back", version=current_version)
+            return
+
+        # `to_roll_back` is the top (highest-version) contiguous suffix of `applied`. Walk it
+        # from the highest version down, and after rolling back applied[idx], the new current
+        # version is its predecessor in the *full* applied list — applied[idx - 1], or 0 once
+        # nothing is left.
+        idx = len(applied) - 1
+        for version in reversed(to_roll_back):
+            assert applied[idx] == version
+            migration = migrations[version]
+            if migration.down_sql is None:
+                raise NoMigrationsFoundError(f"missing *.down.sql for version {version}")
+            idx -= 1
+            new_version = applied[idx] if idx >= 0 else 0
+            await self._apply_sql(
+                conn, migration.down_sql, failing_version=version, success_version=new_version
+            )
+            if self._logger is not None:
+                self._logger.info("rolled back migration", version=version, direction="down")

pgcommon/pool.py

@@ -0,0 +1,346 @@
+"""The connection pool — the only sanctioned entry point to PostgreSQL.
+
+Wraps `asyncpg.Pool`. There is intentionally no public API that exposes the underlying
+`asyncpg.Pool` or a raw `asyncpg.Connection` outside a `with_conn` / `run_in_tx` callback —
+bypassing these helpers silently disables GUC injection, retry logic, and observability.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import random
+import time
+from collections.abc import AsyncIterator, Awaitable, Callable
+from contextlib import asynccontextmanager
+from dataclasses import dataclass
+from datetime import timedelta
+from typing import TypeVar
+
+import asyncpg
+
+from pgcommon import metrics as metrics_module
+from pgcommon._internal.dsn import mask_dsn
+from pgcommon._internal.query_event import QueryEvent
+from pgcommon.config import Config
+from pgcommon.errors import PoolClosedError
+from pgcommon.guc import _inject_gucs_local, inject_gucs
+from pgcommon.metrics import OTelQueryRecorder
+from pgcommon.tracer import SlowQueryTracer
+
+__all__ = ["HealthStatus", "Pool", "Stat"]
+
+T = TypeVar("T")
+
+_DEFAULT_HEALTH_PING_TIMEOUT = 5.0
+
+
+@dataclass(frozen=True, slots=True)
+class Stat:
+    """Snapshot of pool sizing and acquire counters — see `Pool.stat()`."""
+
+    size: int
+    idle_size: int
+    min_size: int
+    max_size: int
+    acquired_size: int
+    acquire_count: int
+    acquire_duration: timedelta
+    empty_acquire_count: int
+    canceled_acquire_count: int
+
+
+@dataclass(frozen=True, slots=True)
+class HealthStatus:
+    """Point-in-time health snapshot returned by `Pool.health()` — call per-request, not once."""
+
+    healthy: bool
+    total_conns: int
+    idle_conns: int
+    acquired_conns: int
+    max_conns: int
+    utilization: float
+
+
+class Pool:
+    """Construct via `await Pool.create(config)` — never call `Pool()` directly."""
+
+    def __init__(self, *, config: Config, asyncpg_pool: asyncpg.Pool) -> None:
+        self._config = config
+        self._pool = asyncpg_pool
+        self._closed = False
+        self._acquire_count = 0
+        self._empty_acquire_count = 0
+        self._canceled_acquire_count = 0
+        self._total_acquire_seconds = 0.0
+        self._background_tasks: list[asyncio.Task[None]] = []
+
+    @classmethod
+    async def create(cls, config: Config) -> Pool:
+        """Validate `config`, open the pool, and ping it once so misconfiguration fails fast."""
+        config.validate()
+
+        slow_query_tracer: SlowQueryTracer | None = None
+        if config.logger is not None and config.slow_query_threshold > timedelta(0):
+            slow_query_tracer = SlowQueryTracer(
+                config.slow_query_threshold,
+                config.logger,
+                log_tenant_id=config.log_tenant_id,
+                allow_full_statements=config.allow_full_statements,
+                guc_provider=config.guc_provider,
+                on_slow_query=lambda: _inc_slow_query_total(),
+            )
+
+        otel_recorder: OTelQueryRecorder | None = None
+        if config.tracer is not None:
+            otel_recorder = OTelQueryRecorder(
+                config.tracer,
+                allow_full_statements=config.allow_full_statements,
+                metrics=metrics_module.current(),
+            )
+
+        def _inc_slow_query_total() -> None:
+            m = metrics_module.current()
+            if m is not None:
+                m.slow_query_total.inc()
+
+        def _on_query(record: QueryEvent) -> None:
+            if slow_query_tracer is not None:
+                slow_query_tracer(record)
+            if otel_recorder is not None:
+                otel_recorder(record)
+
+        has_query_hooks = slow_query_tracer is not None or otel_recorder is not None
+
+        async def _init(conn: asyncpg.Connection) -> None:
+            if has_query_hooks:
+                conn.add_query_logger(_on_query)
+
+        async def _setup(conn: asyncpg.pool.PoolConnectionProxy) -> None:
+            if config.guc_provider is not None and not config.pgbouncer_mode:
+                guc = config.guc_provider()
+                if guc is not None:
+                    await inject_gucs(conn, guc)
+
+        connect_kwargs: dict[str, object] = {}
+        if config.pgbouncer_mode:
+            # PgBouncer transaction pooling can hand back a different backend session on
+            # every acquire; asyncpg's client-side statement cache must be disabled or it
+            # may reuse a prepared statement name the new backend never actually prepared.
+            connect_kwargs["statement_cache_size"] = 0
+
+        try:
+            asyncpg_pool = await asyncpg.create_pool(
+                dsn=config.dsn,
+                min_size=config.min_conns,
+                max_size=config.max_conns,
+                max_inactive_connection_lifetime=config.max_conn_idle_time.total_seconds(),
+                init=_init,
+                setup=_setup,
+                **connect_kwargs,
+            )
+        except Exception as exc:
+            raise ConnectionError(
+                f"failed to create pool for dsn={mask_dsn(config.dsn)}: {mask_dsn(str(exc))}"
+            ) from exc
+
+        if asyncpg_pool is None:  # pragma: no cover - asyncpg always returns a Pool here
+            raise ConnectionError("asyncpg.create_pool unexpectedly returned None")
+
+        pool = cls(config=config, asyncpg_pool=asyncpg_pool)
+
+        try:
+            await pool.ping()
+        except Exception as exc:
+            await asyncpg_pool.close()
+            raise ConnectionError(
+                f"initial ping failed for dsn={mask_dsn(config.dsn)}: {mask_dsn(str(exc))}"
+            ) from exc
+
+        pool._start_background_tasks()
+        return pool
+
+    @asynccontextmanager
+    async def _acquire(self) -> AsyncIterator[asyncpg.pool.PoolConnectionProxy]:
+        """Acquire a raw connection with acquire counters/metrics recorded, releasing on exit.
+
+        Shared primitive behind `with_conn` and `pgcommon.tx.run_in_tx` — deliberately does
+        *not* perform any GUC injection itself, since the two callers scope GUCs differently
+        (an implicit wrapping transaction for `with_conn`, the caller-controlled transaction
+        for `run_in_tx`).
+        """
+        if self._closed:
+            raise PoolClosedError("pool is closed")
+
+        was_empty = self._pool.get_idle_size() == 0
+        start = time.monotonic()
+        cm = self._pool.acquire()
+        try:
+            conn = await cm.__aenter__()
+        except asyncio.CancelledError:
+            self._canceled_acquire_count += 1
+            self._record_acquire_metric(start, status="error")
+            raise
+        except BaseException:
+            self._record_acquire_metric(start, status="error")
+            raise
+
+        self._record_acquire_success(start, was_empty)
+        try:
+            yield conn
+        finally:
+            # PoolAcquireContext.__aexit__ ignores exception info and only releases the
+            # connection, so it's always safe to call unconditionally here.
+            await cm.__aexit__(None, None, None)
+
+    async def with_conn(self, fn: Callable[[asyncpg.pool.PoolConnectionProxy], Awaitable[T]]) -> T:
+        """Acquire a connection, inject GUCs, run `fn`, then release — the only safe entry point.
+
+        In PgBouncer mode, when a `GUCSet` is active, `fn` runs inside an implicit wrapping
+        transaction used to scope the injected GUCs as `is_local=true`. Do not call
+        `conn.transaction()` yourself inside `fn` in that case — see ARCHITECTURE.md.
+        """
+        async with self._acquire() as conn:
+            if self._config.pgbouncer_mode and self._config.guc_provider is not None:
+                guc = self._config.guc_provider()
+                if guc is not None and not guc.is_zero():
+                    guc.validate()
+                    async with conn.transaction():
+                        await _inject_gucs_local(conn, guc)
+                        return await fn(conn)
+            return await fn(conn)
+
+    async def run_in_tx(
+        self,
+        fn: Callable[[asyncpg.pool.PoolConnectionProxy], Awaitable[T]],
+        *,
+        isolation: str | None = None,
+        read_only: bool = False,
+    ) -> T:
+        """Convenience method form of `pgcommon.tx.run_in_tx(pool, fn, ...)`."""
+        from pgcommon.tx import run_in_tx as _run_in_tx
+
+        return await _run_in_tx(self, fn, isolation=isolation, read_only=read_only)
+
+    def _record_acquire_metric(self, start: float, *, status: str) -> None:
+        elapsed = time.monotonic() - start
+        self._total_acquire_seconds += elapsed
+        m = metrics_module.current()
+        if m is not None:
+            m.pool_acquire_total.labels(status=status).inc()
+            m.pool_acquire_duration.observe(elapsed)
+
+    def _record_acquire_success(self, start: float, was_empty: bool) -> None:
+        self._acquire_count += 1
+        if was_empty:
+            self._empty_acquire_count += 1
+        self._record_acquire_metric(start, status="success")
+
+    async def ping(self) -> None:
+        """Verify database connectivity. Raises `PoolClosedError` if the pool has been closed."""
+        if self._closed:
+            raise PoolClosedError("pool is closed")
+        async with self._pool.acquire() as conn:
+            await conn.execute("SELECT 1")
+
+    async def health(self, *, timeout: float = _DEFAULT_HEALTH_PING_TIMEOUT) -> HealthStatus:
+        """Return a point-in-time `HealthStatus`. Never raises — ping failures set `healthy=False`."""
+        stat = self.stat()
+        try:
+            await asyncio.wait_for(self.ping(), timeout=timeout)
+            healthy = True
+        except Exception:
+            healthy = False
+
+        max_conns = self._config.max_conns
+        utilization = (stat.acquired_size / max_conns) if max_conns > 0 else 0.0
+        return HealthStatus(
+            healthy=healthy,
+            total_conns=stat.size,
+            idle_conns=stat.idle_size,
+            acquired_conns=stat.acquired_size,
+            max_conns=max_conns,
+            utilization=utilization,
+        )
+
+    def stat(self) -> Stat:
+        """Return a snapshot of pool sizing and acquire counters."""
+        size = self._pool.get_size()
+        idle_size = self._pool.get_idle_size()
+        return Stat(
+            size=size,
+            idle_size=idle_size,
+            min_size=self._pool.get_min_size(),
+            max_size=self._pool.get_max_size(),
+            acquired_size=size - idle_size,
+            acquire_count=self._acquire_count,
+            acquire_duration=timedelta(seconds=self._total_acquire_seconds),
+            empty_acquire_count=self._empty_acquire_count,
+            canceled_acquire_count=self._canceled_acquire_count,
+        )
+
+    def close(self) -> None:
+        """Immediately invalidate the pool. Prefer `drain_and_close` during graceful shutdown."""
+        self._closed = True
+        self._cancel_background_tasks()
+        self._pool.terminate()
+
+    async def drain_and_close(self, *, timeout: float = 30.0) -> None:
+        """Mark the pool closed, wait for in-flight callbacks to finish, then force-close.
+
+        New `with_conn` / `ping` / `health` calls raise `PoolClosedError` immediately. Already
+        in-flight callbacks are allowed to finish; if they have not all finished by `timeout`
+        seconds, the pool is force-closed.
+        """
+        self._closed = True
+        self._cancel_background_tasks()
+        try:
+            await asyncio.wait_for(self._pool.close(), timeout=timeout)
+        except TimeoutError:
+            self._pool.terminate()
+
+    def _start_background_tasks(self) -> None:
+        if self._config.max_conn_lifetime > timedelta(0):
+            self._background_tasks.append(
+                asyncio.ensure_future(self._recycle_connections_periodically())
+            )
+        if self._config.health_check_period > timedelta(0):
+            self._background_tasks.append(asyncio.ensure_future(self._probe_health_periodically()))
+
+    def _cancel_background_tasks(self) -> None:
+        for task in self._background_tasks:
+            task.cancel()
+        self._background_tasks.clear()
+
+    async def _recycle_connections_periodically(self) -> None:
+        """Recycle the whole pool generation every `max_conn_lifetime` (+ jitter).
+
+        This is a deliberate simplification from pgx's per-connection independent aging:
+        asyncpg has no public hook to force-expire a single physical connection without
+        surfacing an error to whichever caller currently holds it. `expire_connections()`
+        instead bumps a pool-wide generation counter so every connection is transparently
+        replaced the next time it is idle and reacquired. The jitter still serves its
+        purpose — staggering *when* different service instances recycle their pools — even
+        though it no longer staggers individual connections within one pool.
+        """
+        base = self._config.max_conn_lifetime.total_seconds()
+        jitter = self._config.max_conn_lifetime_jitter.total_seconds()
+        try:
+            while True:
+                delay = base + (random.uniform(0, jitter) if jitter > 0 else 0.0)
+                await asyncio.sleep(delay)
+                await self._pool.expire_connections()
+        except asyncio.CancelledError:
+            return
+
+    async def _probe_health_periodically(self) -> None:
+        interval = self._config.health_check_period.total_seconds()
+        try:
+            while True:
+                await asyncio.sleep(interval)
+                try:
+                    await self.ping()
+                except Exception as exc:
+                    if self._config.logger is not None:
+                        self._config.logger.warn("pool health check failed", error=str(exc))
+        except asyncio.CancelledError:
+            return