pgcommon 0.1.0__py3-none-any.whl

pgcommon/retry.py

@@ -0,0 +1,121 @@
+"""Retry helpers for transient serialization failures and deadlocks.
+
+⚠️ **Retry does NOT make an operation safe — it re-runs it.** These helpers re-execute `fn`
+from the top, up to `max_attempts` times. `fn` must be idempotent: the transaction rolls back
+between attempts, undoing *database* writes, but it cannot undo a side effect that already
+happened outside the transaction (an email send, a payment charge, a webhook POST). Do all
+non-database I/O outside the transaction, or make it idempotent with a key. See the README's
+"Retry does NOT make an operation safe" section before using these on a write path.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import random
+from collections.abc import Awaitable, Callable
+from dataclasses import dataclass
+from datetime import timedelta
+from typing import TYPE_CHECKING
+
+import asyncpg
+
+from pgcommon import metrics as metrics_module
+from pgcommon.errors import is_deadlock, is_serialization_failure
+from pgcommon.tx import run_in_tx
+
+if TYPE_CHECKING:
+    from pgcommon.pool import Pool
+
+__all__ = ["RetryOptions", "run_in_tx_with_retry", "run_in_tx_with_retry_opts"]
+
+
+@dataclass(frozen=True, slots=True)
+class RetryOptions:
+    """Configures the backoff schedule for `run_in_tx_with_retry_opts`.
+
+    delay = min(initial_wait * multiplier**(attempt-1), max_wait)   # when max_wait > 0
+    delay ± jitter_fraction * delay                                 # randomized ± spread
+    """
+
+    max_attempts: int = 3
+    initial_wait: timedelta = timedelta(0)
+    max_wait: timedelta = timedelta(0)
+    multiplier: float = 1.0
+    jitter_fraction: float = 0.0
+
+
+def _retry_reason(exc: BaseException) -> str | None:
+    if is_deadlock(exc):
+        return "deadlock"
+    if is_serialization_failure(exc):
+        return "serialization"
+    return None
+
+
+def _compute_delay_seconds(opts: RetryOptions, attempt: int) -> float:
+    delay = opts.initial_wait.total_seconds() * (opts.multiplier ** (attempt - 1))
+    max_wait_seconds = opts.max_wait.total_seconds()
+    if max_wait_seconds > 0:
+        delay = min(delay, max_wait_seconds)
+    if opts.jitter_fraction > 0:
+        delay += random.uniform(-opts.jitter_fraction, opts.jitter_fraction) * delay
+    return max(delay, 0.0)
+
+
+async def run_in_tx_with_retry[T](
+    pool: Pool,
+    fn: Callable[[asyncpg.pool.PoolConnectionProxy], Awaitable[T]],
+    max_attempts: int,
+    *,
+    isolation: str | None = None,
+    read_only: bool = False,
+) -> T:
+    """Retry on deadlock (40P01) / serialization failure (40001), up to `max_attempts`.
+
+    Retries immediately, with no delay between attempts. Non-retryable errors are returned to
+    the caller unchanged on the first occurrence.
+    """
+    return await run_in_tx_with_retry_opts(
+        pool,
+        fn,
+        RetryOptions(max_attempts=max_attempts),
+        isolation=isolation,
+        read_only=read_only,
+    )
+
+
+async def run_in_tx_with_retry_opts[T](
+    pool: Pool,
+    fn: Callable[[asyncpg.pool.PoolConnectionProxy], Awaitable[T]],
+    opts: RetryOptions,
+    *,
+    isolation: str | None = None,
+    read_only: bool = False,
+) -> T:
+    """Like `run_in_tx_with_retry`, with configurable exponential backoff and jitter.
+
+    `asyncio.CancelledError` is never treated as retryable — it propagates immediately,
+    whether raised during the transaction attempt or during the backoff sleep.
+    """
+    attempt = 1
+    while True:
+        try:
+            return await run_in_tx(pool, fn, isolation=isolation, read_only=read_only)
+        except Exception as exc:
+            reason = _retry_reason(exc)
+            if reason is None:
+                raise
+
+            m = metrics_module.current()
+            if m is not None:
+                m.retry_total.labels(reason=reason).inc()
+
+            if attempt >= opts.max_attempts:
+                if m is not None:
+                    m.retry_exhausted_total.labels(reason=reason).inc()
+                raise
+
+            delay = _compute_delay_seconds(opts, attempt)
+            if delay > 0:
+                await asyncio.sleep(delay)
+            attempt += 1

pgcommon/tracer.py

@@ -0,0 +1,85 @@
+"""Slow-query detection and structured logging.
+
+Mirrors the Go repo's `SlowQueryTracer`. pgx exposes a `QueryTracer` interface with paired
+`TraceQueryStart` / `TraceQueryEnd` hooks; asyncpg's closest analog is
+`Connection.add_query_logger`, a single post-query callback carrying `elapsed` and
+`exception`. `SlowQueryTracer` is written as a callable so it can be registered directly.
+"""
+
+from __future__ import annotations
+
+import time
+from collections.abc import Callable
+from datetime import timedelta
+
+from pgcommon._internal.query_event import QueryEvent
+from pgcommon._internal.sql_verb import sql_operation
+from pgcommon.guc import GUCSet
+from pgcommon.logger import Logger
+
+__all__ = ["SlowQueryTracer"]
+
+_RATE_LIMIT_SECONDS = 1.0
+
+
+class SlowQueryTracer:
+    """Logs `WARN` for any query exceeding `threshold`, rate-limited to 1 log/second.
+
+    The `pgcommon_slow_query_total` counter (wired via `on_slow_query`) fires on every
+    detection regardless of the log rate limit — only the log line itself is throttled.
+    """
+
+    def __init__(
+        self,
+        threshold: timedelta,
+        logger: Logger,
+        *,
+        log_tenant_id: bool = False,
+        allow_full_statements: bool = False,
+        guc_provider: Callable[[], GUCSet | None] | None = None,
+        on_slow_query: Callable[[], None] | None = None,
+    ) -> None:
+        self._threshold_seconds = threshold.total_seconds()
+        self._logger = logger
+        self._log_tenant_id = log_tenant_id
+        self._allow_full_statements = allow_full_statements
+        self._guc_provider = guc_provider
+        self._on_slow_query = on_slow_query
+        self._last_log_monotonic = float("-inf")
+
+    def set_log_tenant_id(self, value: bool) -> None:
+        """Opt in (or back out) of including `tenant_id` in slow-query log entries."""
+        self._log_tenant_id = value
+
+    def set_allow_full_statements(self, value: bool) -> None:
+        """Opt in (or back out) of logging full SQL text instead of just the verb."""
+        self._allow_full_statements = value
+
+    def set_slow_query_hook(self, hook: Callable[[], None] | None) -> None:
+        """Set (or clear) the callback invoked on every slow-query detection."""
+        self._on_slow_query = hook
+
+    def __call__(self, record: QueryEvent) -> None:
+        """Register this as an asyncpg query logger: `conn.add_query_logger(tracer)`."""
+        if record.elapsed < self._threshold_seconds:
+            return
+
+        if self._on_slow_query is not None:
+            self._on_slow_query()
+
+        now = time.monotonic()
+        if now - self._last_log_monotonic < _RATE_LIMIT_SECONDS:
+            return
+        self._last_log_monotonic = now
+
+        fields: dict[str, object] = {
+            "query": record.query if self._allow_full_statements else sql_operation(record.query),
+            "args_count": len(record.args) if record.args is not None else 0,
+            "duration_ms": round(record.elapsed * 1000, 3),
+        }
+        if self._log_tenant_id and self._guc_provider is not None:
+            guc = self._guc_provider()
+            if guc is not None and guc.tenant_id:
+                fields["tenant_id"] = guc.tenant_id
+
+        self._logger.warn("slow query", **fields)

pgcommon/tx.py

@@ -0,0 +1,98 @@
+"""Transaction helpers: `run_in_tx` and `run_in_savepoint`.
+
+`run_in_tx` commits on success and rolls back on exception, including `asyncio.CancelledError`
+— handled by asyncpg's own `conn.transaction()` context manager, whose `__aexit__` sends
+`ROLLBACK` when any exception (cancellation included) propagates out of the `async with` block.
+
+This is a deliberately different mechanism from the Go repo's "rollback uses
+`context.Background()`" trick, not just a different spelling of the same one — see
+ARCHITECTURE.md's Cancellation section for why Go's workaround doesn't have a direct Python
+analog: Go's `context.Context` is checked by the driver *before* it will even attempt an I/O
+call, so a already-cancelled context can suppress the ROLLBACK from being sent at all. asyncio
+cancellation instead interrupts exactly one suspended `await`; once that single
+`CancelledError` has been delivered and is being handled in cleanup code, further `await`s
+(including the one sending `ROLLBACK`) run to completion normally unless cancelled a second
+time. That is why no `asyncio.shield()` or "background context" equivalent is needed here.
+"""
+
+from __future__ import annotations
+
+import re
+from collections.abc import Awaitable, Callable
+from typing import TYPE_CHECKING
+
+import asyncpg
+
+from pgcommon.errors import InvalidSavepointNameError
+from pgcommon.guc import _inject_gucs_local
+
+if TYPE_CHECKING:
+    from pgcommon.pool import Pool
+
+__all__ = ["run_in_savepoint", "run_in_tx"]
+
+_SAVEPOINT_NAME_RE = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$")
+
+
+async def run_in_tx[T](
+    pool: Pool,
+    fn: Callable[[asyncpg.pool.PoolConnectionProxy], Awaitable[T]],
+    *,
+    isolation: str | None = None,
+    read_only: bool = False,
+) -> T:
+    """Run `fn` inside a transaction: commit on success, rollback on exception or cancellation.
+
+    `isolation` is passed straight through to `asyncpg.Connection.transaction()` — one of
+    `"read_committed"`, `"repeatable_read"`, `"serializable"`, or `None` for the session
+    default.
+
+    In PgBouncer mode, when a `GUCSet` is active on the current context, GUCs are re-injected
+    (`is_local=true`) inside this same transaction before `fn` runs, so they are automatically
+    cleared on commit/rollback rather than leaking into the next client that gets this backend
+    connection from PgBouncer.
+    """
+    async with pool._acquire() as conn:
+        config = pool._config
+        guc = None
+        if config.pgbouncer_mode and config.guc_provider is not None:
+            candidate = config.guc_provider()
+            if candidate is not None and not candidate.is_zero():
+                candidate.validate()
+                guc = candidate
+
+        tx_kwargs: dict[str, object] = {}
+        if isolation is not None:
+            tx_kwargs["isolation"] = isolation
+        if read_only:
+            tx_kwargs["readonly"] = True
+
+        async with conn.transaction(**tx_kwargs):
+            if guc is not None:
+                await _inject_gucs_local(conn, guc)
+            return await fn(conn)
+
+
+async def run_in_savepoint[T](
+    conn: asyncpg.pool.PoolConnectionProxy,
+    name: str,
+    fn: Callable[[asyncpg.pool.PoolConnectionProxy], Awaitable[T]],
+) -> T:
+    """Run `fn` inside a named `SAVEPOINT` nested in the caller's already-open transaction.
+
+    Safe to call from inside a `run_in_tx` callback. On success: `RELEASE SAVEPOINT`. On
+    exception: `ROLLBACK TO SAVEPOINT` — the outer transaction is unaffected and the caller may
+    continue using it. `name` must match `^[A-Za-z_][A-Za-z0-9_]*$`.
+    """
+    if not _SAVEPOINT_NAME_RE.match(name):
+        raise InvalidSavepointNameError(f"invalid savepoint name: {name!r}")
+
+    await conn.execute(f"SAVEPOINT {name}")
+    try:
+        result = await fn(conn)
+    except BaseException:
+        await conn.execute(f"ROLLBACK TO SAVEPOINT {name}")
+        raise
+    else:
+        await conn.execute(f"RELEASE SAVEPOINT {name}")
+        return result