pgcommon 0.1.0__py3-none-any.whl

pgcommon/__init__.py

@@ -0,0 +1,91 @@
+"""pgcommon — opinionated async PostgreSQL access layer for platform services.
+
+Public re-exports only. Anything under `pgcommon._internal` is private and must never be
+imported by consuming services.
+"""
+
+from __future__ import annotations
+
+from pgcommon.config import Config, ConfigWarning, config_from_env
+from pgcommon.errors import (
+    DSNRequiredError,
+    EmptyTenantRoleError,
+    InconsistentGUCSetError,
+    InvalidMaxConnsError,
+    InvalidMinConnsError,
+    InvalidSavepointNameError,
+    InvalidTenantRoleError,
+    MigrationDirtyError,
+    MinConnsExceedsMaxError,
+    NoMigrationsFoundError,
+    PgCommonError,
+    PoolClosedError,
+    constraint_name,
+    is_check_violation,
+    is_deadlock,
+    is_foreign_key_violation,
+    is_not_null_violation,
+    is_serialization_failure,
+    is_unique_violation,
+)
+from pgcommon.guc import (
+    GUCSet,
+    guc_set_from_context,
+    inject_gucs,
+    with_guc_set,
+    with_validated_guc_set,
+)
+from pgcommon.logger import Logger, StdlibLogger, StructlogLogger
+from pgcommon.metrics import Metrics, OTelQueryRecorder
+from pgcommon.metrics import init as pgmetrics_init
+from pgcommon.metrics import init_with_registry as pgmetrics_init_with_registry
+from pgcommon.pool import HealthStatus, Pool, Stat
+from pgcommon.retry import RetryOptions, run_in_tx_with_retry, run_in_tx_with_retry_opts
+from pgcommon.tracer import SlowQueryTracer
+from pgcommon.tx import run_in_savepoint, run_in_tx
+
+__all__ = [
+    "Config",
+    "ConfigWarning",
+    "DSNRequiredError",
+    "EmptyTenantRoleError",
+    "GUCSet",
+    "HealthStatus",
+    "InconsistentGUCSetError",
+    "InvalidMaxConnsError",
+    "InvalidMinConnsError",
+    "InvalidSavepointNameError",
+    "InvalidTenantRoleError",
+    "Logger",
+    "Metrics",
+    "MigrationDirtyError",
+    "MinConnsExceedsMaxError",
+    "NoMigrationsFoundError",
+    "OTelQueryRecorder",
+    "PgCommonError",
+    "Pool",
+    "PoolClosedError",
+    "RetryOptions",
+    "SlowQueryTracer",
+    "Stat",
+    "StdlibLogger",
+    "StructlogLogger",
+    "config_from_env",
+    "constraint_name",
+    "guc_set_from_context",
+    "inject_gucs",
+    "is_check_violation",
+    "is_deadlock",
+    "is_foreign_key_violation",
+    "is_not_null_violation",
+    "is_serialization_failure",
+    "is_unique_violation",
+    "pgmetrics_init",
+    "pgmetrics_init_with_registry",
+    "run_in_savepoint",
+    "run_in_tx",
+    "run_in_tx_with_retry",
+    "run_in_tx_with_retry_opts",
+    "with_guc_set",
+    "with_validated_guc_set",
+]

pgcommon/_internal/__init__.py

@@ -0,0 +1 @@
+"""Private helpers for pgcommon. Not part of the public API — do not import from outside this package."""

pgcommon/_internal/dsn.py

@@ -0,0 +1,15 @@
+from __future__ import annotations
+
+import re
+
+# Matches the password segment of a URI-style DSN: scheme://user:PASSWORD@host
+_URI_PASSWORD = re.compile(r"(://[^:/?#@]*:)([^@/?#]*)(@)")
+# Matches ` password=... ` or `password='...'` in keyword/value DSNs.
+_KV_PASSWORD = re.compile(r"(?i)(\bpassword\s*=\s*)('(?:[^'\\]|\\.)*'|\S+)")
+
+
+def mask_dsn(dsn: str) -> str:
+    """Replace the password segment of a DSN with '***' so it is safe to include in error messages."""
+    masked = _URI_PASSWORD.sub(r"\1***\3", dsn)
+    masked = _KV_PASSWORD.sub(r"\1***", masked)
+    return masked

pgcommon/_internal/duration.py

@@ -0,0 +1,38 @@
+from __future__ import annotations
+
+import re
+from datetime import timedelta
+
+_TOKEN = re.compile(r"([0-9]*\.?[0-9]+)(ns|us|µs|ms|s|m|h)")
+_UNIT_SECONDS: dict[str, float] = {
+    "ns": 1e-9,
+    "us": 1e-6,
+    "µs": 1e-6,
+    "ms": 1e-3,
+    "s": 1.0,
+    "m": 60.0,
+    "h": 3600.0,
+}
+
+
+def parse_duration(value: str) -> timedelta:
+    """Parse a Go `time.ParseDuration`-style string ("200ms", "1h30m") into a `timedelta`."""
+    text = value.strip()
+    if not text:
+        raise ValueError("empty duration string")
+
+    pos = 0
+    total_seconds = 0.0
+    matched_any = False
+    for m in _TOKEN.finditer(text):
+        if m.start() != pos:
+            raise ValueError(f"invalid duration {value!r}")
+        pos = m.end()
+        num, unit = m.groups()
+        total_seconds += float(num) * _UNIT_SECONDS[unit]
+        matched_any = True
+
+    if not matched_any or pos != len(text):
+        raise ValueError(f"invalid duration {value!r}")
+
+    return timedelta(seconds=total_seconds)

pgcommon/_internal/query_event.py

@@ -0,0 +1,22 @@
+from __future__ import annotations
+
+from collections.abc import Sequence
+from typing import Protocol, runtime_checkable
+
+
+@runtime_checkable
+class QueryEvent(Protocol):
+    """Structural shape of asyncpg's `Connection.add_query_logger` callback record.
+
+    asyncpg has no pre/post-query hook pair like pgx's `QueryTracer` — instead a single
+    callback fires after each query completes, carrying elapsed time and any exception. Spans
+    are reconstructed after the fact using `elapsed` to back-date `start_time`; see
+    `pgcommon.metrics.OTelQueryRecorder`.
+    """
+
+    query: str
+    args: Sequence[object] | None
+    elapsed: float
+    exception: BaseException | None
+    conn_addr: tuple[str, int] | str | None
+    conn_params: object

pgcommon/_internal/sql_verb.py

@@ -0,0 +1,36 @@
+from __future__ import annotations
+
+_KNOWN_VERBS = frozenset(
+    {
+        "SELECT",
+        "INSERT",
+        "UPDATE",
+        "DELETE",
+        "BEGIN",
+        "COMMIT",
+        "ROLLBACK",
+        "SAVEPOINT",
+        "RELEASE",
+        "CREATE",
+        "ALTER",
+        "DROP",
+        "TRUNCATE",
+        "WITH",
+        "EXPLAIN",
+        "COPY",
+        "SET",
+        "GRANT",
+        "REVOKE",
+    }
+)
+
+
+def sql_operation(sql: str) -> str:
+    """Extract the leading SQL verb from a statement, for use as a bounded-cardinality metric/span label.
+
+    Never returns raw SQL — this is what keeps `pgcommon_query_total{operation}` and
+    `db.statement` (when full statements are not allowed) free of PII and unbounded cardinality.
+    """
+    stripped = sql.lstrip()
+    first_word = stripped.split(None, 1)[0].upper() if stripped else ""
+    return first_word if first_word in _KNOWN_VERBS else "other"

pgcommon/_internal/sqlstate.py

@@ -0,0 +1,20 @@
+from __future__ import annotations
+
+UNIQUE_VIOLATION = "23505"
+FOREIGN_KEY_VIOLATION = "23503"
+NOT_NULL_VIOLATION = "23502"
+CHECK_VIOLATION = "23514"
+DEADLOCK_DETECTED = "40P01"
+SERIALIZATION_FAILURE = "40001"
+
+RETRYABLE = frozenset({DEADLOCK_DETECTED, SERIALIZATION_FAILURE})
+
+
+def sqlstate_of(exc: BaseException) -> str | None:
+    """Return the SQLSTATE code carried by a PostgresError-like exception, if any."""
+    return getattr(exc, "sqlstate", None)
+
+
+def constraint_name_of(exc: BaseException) -> str:
+    """Return the offending constraint name carried by a PostgresError-like exception, or ""."""
+    return getattr(exc, "constraint_name", None) or ""

pgcommon/config.py

@@ -0,0 +1,222 @@
+"""Pool configuration: the `Config` dataclass, `config_from_env()`, and `ConfigWarning`.
+
+Pure and I/O-free except for `config_from_env`, which reads `os.environ` — mirroring the Go
+repo's `internal/config.Load()`. No network or database access happens here; `Pool.create`
+is what fails fast against a live connection.
+"""
+
+from __future__ import annotations
+
+import os
+from collections.abc import Callable
+from dataclasses import dataclass
+from datetime import timedelta
+from typing import TYPE_CHECKING
+
+from pgcommon._internal.duration import parse_duration
+from pgcommon.errors import (
+    DSNRequiredError,
+    InvalidMaxConnsError,
+    InvalidMinConnsError,
+    MinConnsExceedsMaxError,
+)
+from pgcommon.guc import GUCSet
+from pgcommon.logger import Logger
+
+if TYPE_CHECKING:
+    from opentelemetry.trace import Tracer
+
+__all__ = [
+    "Config",
+    "ConfigWarning",
+    "config_from_env",
+]
+
+_DEFAULT_MAX_CONNS = 10
+_DEFAULT_MIN_CONNS = 2
+_DEFAULT_MAX_CONN_LIFETIME = timedelta(hours=1)
+_DEFAULT_MAX_CONN_LIFETIME_JITTER = timedelta(0)
+_DEFAULT_MAX_CONN_IDLE_TIME = timedelta(minutes=30)
+_DEFAULT_HEALTH_CHECK_PERIOD = timedelta(minutes=1)
+_DEFAULT_SLOW_QUERY_THRESHOLD = timedelta(milliseconds=200)
+_INSECURE_SSLMODES = frozenset({"disable", "allow", "prefer"})
+
+
+@dataclass(slots=True)
+class Config:
+    """Pool configuration. `dsn` is the only required field."""
+
+    dsn: str
+    max_conns: int = _DEFAULT_MAX_CONNS
+    min_conns: int = _DEFAULT_MIN_CONNS
+    max_conn_lifetime: timedelta = _DEFAULT_MAX_CONN_LIFETIME
+    max_conn_lifetime_jitter: timedelta = _DEFAULT_MAX_CONN_LIFETIME_JITTER
+    max_conn_idle_time: timedelta = _DEFAULT_MAX_CONN_IDLE_TIME
+    health_check_period: timedelta = _DEFAULT_HEALTH_CHECK_PERIOD
+    slow_query_threshold: timedelta = _DEFAULT_SLOW_QUERY_THRESHOLD
+    logger: Logger | None = None
+    tracer: Tracer | None = None
+    guc_provider: Callable[[], GUCSet | None] | None = None
+    pgbouncer_mode: bool = False
+    log_tenant_id: bool = False
+    allow_full_statements: bool = False
+
+    def validate(self) -> None:
+        """Raise `DSNRequiredError` / `InvalidMaxConnsError` / `InvalidMinConnsError` /
+        `MinConnsExceedsMaxError`.
+
+        Called by `Pool.create` before any connection is attempted — misconfiguration is a
+        startup error, never a runtime surprise.
+        """
+        if not self.dsn:
+            raise DSNRequiredError("Config.dsn is required")
+        if self.max_conns < 0:
+            raise InvalidMaxConnsError(f"max_conns must be >= 0, got {self.max_conns}")
+        if self.min_conns < 0:
+            raise InvalidMinConnsError(f"min_conns must be >= 0, got {self.min_conns}")
+        if self.min_conns > self.max_conns:
+            raise MinConnsExceedsMaxError(
+                f"min_conns ({self.min_conns}) exceeds max_conns ({self.max_conns})"
+            )
+
+
+@dataclass(frozen=True, slots=True)
+class ConfigWarning:
+    """A non-fatal configuration issue surfaced by `config_from_env` — log these at startup."""
+
+    key: str
+    reason: str
+
+
+def _parse_positive_int(env_var: str, default: int) -> tuple[int, ConfigWarning | None]:
+    raw = os.environ.get(env_var)
+    if raw is None:
+        return default, None
+    try:
+        parsed = int(raw)
+    except ValueError:
+        return default, ConfigWarning(
+            key=env_var, reason=f"invalid integer {raw!r}, using default {default}"
+        )
+    if parsed < 0:
+        return default, ConfigWarning(
+            key=env_var, reason=f"value {parsed} must be >= 0, using default {default}"
+        )
+    return parsed, None
+
+
+def _parse_duration_with_warning(
+    env_var: str, default: timedelta
+) -> tuple[timedelta, ConfigWarning | None]:
+    raw = os.environ.get(env_var)
+    if raw is None:
+        return default, None
+    try:
+        return parse_duration(raw), None
+    except ValueError:
+        return default, ConfigWarning(
+            key=env_var, reason=f"invalid duration {raw!r}, using default {default}"
+        )
+
+
+def _parse_bool(env_var: str, default: bool) -> bool:
+    raw = os.environ.get(env_var)
+    if raw is None:
+        return default
+    return raw.strip().lower() in ("1", "true", "t", "yes", "y", "on")
+
+
+def _build_dsn() -> tuple[str, ConfigWarning | None]:
+    database_url = os.environ.get("DATABASE_URL")
+    if database_url:
+        return database_url, None
+
+    host = os.environ.get("PG_HOST", "localhost")
+    port = os.environ.get("PG_PORT", "5432")
+    user = os.environ.get("PG_USER", "")
+    password = os.environ.get("PG_PASSWORD", "")
+    dbname = os.environ.get("PG_DBNAME", "")
+    sslmode = os.environ.get("PG_SSLMODE", "require")
+
+    userinfo = f"{user}:{password}@" if user or password else ""
+    dsn = f"postgres://{userinfo}{host}:{port}/{dbname}?sslmode={sslmode}"
+    return dsn, None
+
+
+def _sslmode_of(dsn: str) -> str | None:
+    if "sslmode=" not in dsn:
+        return None
+    tail = dsn.split("sslmode=", 1)[1]
+    for sep in ("&", " "):
+        tail = tail.split(sep, 1)[0]
+    return tail
+
+
+def config_from_env() -> tuple[Config, list[ConfigWarning]]:
+    """Build a `Config` from `DATABASE_URL` / `PG_*` environment variables.
+
+    Mirrors the Go repo's `internal/config.Load()`: invalid integers/durations fall back to
+    their defaults with a `ConfigWarning` rather than raising, and an insecure `sslmode`
+    (`disable`, `allow`, `prefer`) always produces a warning. Log every returned warning at
+    startup so accidental TLS disablement or a typo'd tuning value is visible.
+    """
+    warnings: list[ConfigWarning] = []
+
+    dsn, dsn_warning = _build_dsn()
+    if dsn_warning:
+        warnings.append(dsn_warning)
+
+    max_conns, w = _parse_positive_int("PG_MAX_CONNS", _DEFAULT_MAX_CONNS)
+    if w:
+        warnings.append(w)
+    min_conns, w = _parse_positive_int("PG_MIN_CONNS", _DEFAULT_MIN_CONNS)
+    if w:
+        warnings.append(w)
+    max_conn_lifetime, w = _parse_duration_with_warning(
+        "PG_MAX_CONN_LIFETIME", _DEFAULT_MAX_CONN_LIFETIME
+    )
+    if w:
+        warnings.append(w)
+    max_conn_lifetime_jitter, w = _parse_duration_with_warning(
+        "PG_MAX_CONN_LIFETIME_JITTER", _DEFAULT_MAX_CONN_LIFETIME_JITTER
+    )
+    if w:
+        warnings.append(w)
+    max_conn_idle_time, w = _parse_duration_with_warning(
+        "PG_MAX_CONN_IDLE_TIME", _DEFAULT_MAX_CONN_IDLE_TIME
+    )
+    if w:
+        warnings.append(w)
+    health_check_period, w = _parse_duration_with_warning(
+        "PG_HEALTH_CHECK_PERIOD", _DEFAULT_HEALTH_CHECK_PERIOD
+    )
+    if w:
+        warnings.append(w)
+    slow_query_threshold, w = _parse_duration_with_warning(
+        "PG_SLOW_QUERY_THRESHOLD", _DEFAULT_SLOW_QUERY_THRESHOLD
+    )
+    if w:
+        warnings.append(w)
+    pgbouncer_mode = _parse_bool("PG_BOUNCER_MODE", False)
+
+    sslmode = os.environ.get("PG_SSLMODE") or _sslmode_of(dsn)
+    if sslmode and sslmode.lower() in _INSECURE_SSLMODES:
+        warnings.append(
+            ConfigWarning(
+                key="PG_SSLMODE",
+                reason=f"sslmode={sslmode!r} does not enforce TLS — use 'require' or stronger in production",
+            )
+        )
+
+    cfg = Config(
+        dsn=dsn,
+        max_conns=max_conns,
+        min_conns=min_conns,
+        max_conn_lifetime=max_conn_lifetime,
+        max_conn_lifetime_jitter=max_conn_lifetime_jitter,
+        max_conn_idle_time=max_conn_idle_time,
+        health_check_period=health_check_period,
+        slow_query_threshold=slow_query_threshold,
+        pgbouncer_mode=pgbouncer_mode,
+    )
+    return cfg, warnings

pgcommon/errors.py

@@ -0,0 +1,123 @@
+"""Sentinel exceptions and PostgreSQL error-classification helpers.
+
+Mirrors the Go repo's `domain.Err*` sentinels (check with `isinstance`, the Python
+equivalent of `errors.Is`) and the `Is*Violation` / `ConstraintName` helper functions.
+"""
+
+from __future__ import annotations
+
+from pgcommon._internal.sqlstate import (
+    CHECK_VIOLATION,
+    DEADLOCK_DETECTED,
+    FOREIGN_KEY_VIOLATION,
+    NOT_NULL_VIOLATION,
+    SERIALIZATION_FAILURE,
+    UNIQUE_VIOLATION,
+    constraint_name_of,
+    sqlstate_of,
+)
+
+__all__ = [
+    "DSNRequiredError",
+    "EmptyTenantRoleError",
+    "InconsistentGUCSetError",
+    "InvalidMaxConnsError",
+    "InvalidMinConnsError",
+    "InvalidSavepointNameError",
+    "InvalidTenantRoleError",
+    "MigrationDirtyError",
+    "MinConnsExceedsMaxError",
+    "NoMigrationsFoundError",
+    "PgCommonError",
+    "PoolClosedError",
+    "constraint_name",
+    "is_check_violation",
+    "is_deadlock",
+    "is_foreign_key_violation",
+    "is_not_null_violation",
+    "is_serialization_failure",
+    "is_unique_violation",
+]
+
+
+class PgCommonError(Exception):
+    """Base class for every sentinel error raised by pgcommon."""
+
+
+class DSNRequiredError(PgCommonError):
+    """Raised by `Pool.create` / `Config` validation when no DSN is configured."""
+
+
+class PoolClosedError(PgCommonError):
+    """Raised by `with_conn` / `health` / `ping` after the pool has been closed."""
+
+
+class InvalidMaxConnsError(PgCommonError):
+    """Raised when `Config.max_conns` is negative."""
+
+
+class InvalidMinConnsError(PgCommonError):
+    """Raised when `Config.min_conns` is negative."""
+
+
+class MinConnsExceedsMaxError(PgCommonError):
+    """Raised when `Config.min_conns > Config.max_conns` after defaults are applied."""
+
+
+class InvalidSavepointNameError(PgCommonError):
+    """Raised by `run_in_savepoint` when the name fails `^[A-Za-z_][A-Za-z0-9_]*$`."""
+
+
+class InconsistentGUCSetError(PgCommonError):
+    """Raised by `GUCSet.validate()` when `user_id` is set but `tenant_id` is empty."""
+
+
+class InvalidTenantRoleError(PgCommonError):
+    """Raised by `GUCSet.validate()` when a tenant role contains a comma (the GUC separator)."""
+
+
+class EmptyTenantRoleError(PgCommonError):
+    """Raised by `GUCSet.validate()` when a tenant role is the empty string."""
+
+
+class MigrationDirtyError(PgCommonError):
+    """Raised when the schema is left dirty by a previously failed migration run."""
+
+
+class NoMigrationsFoundError(PgCommonError):
+    """Raised when the migration source directory has no matching `*.sql` files."""
+
+
+def is_unique_violation(exc: BaseException) -> bool:
+    """True when `exc` carries SQLSTATE 23505 (duplicate value in a UNIQUE / PRIMARY KEY column)."""
+    return sqlstate_of(exc) == UNIQUE_VIOLATION
+
+
+def is_foreign_key_violation(exc: BaseException) -> bool:
+    """True when `exc` carries SQLSTATE 23503 (referenced row does not exist)."""
+    return sqlstate_of(exc) == FOREIGN_KEY_VIOLATION
+
+
+def is_not_null_violation(exc: BaseException) -> bool:
+    """True when `exc` carries SQLSTATE 23502 (NULL inserted into a NOT NULL column)."""
+    return sqlstate_of(exc) == NOT_NULL_VIOLATION
+
+
+def is_check_violation(exc: BaseException) -> bool:
+    """True when `exc` carries SQLSTATE 23514 (value failed a CHECK constraint)."""
+    return sqlstate_of(exc) == CHECK_VIOLATION
+
+
+def is_deadlock(exc: BaseException) -> bool:
+    """True when `exc` carries SQLSTATE 40P01 (deadlock detected between transactions)."""
+    return sqlstate_of(exc) == DEADLOCK_DETECTED
+
+
+def is_serialization_failure(exc: BaseException) -> bool:
+    """True when `exc` carries SQLSTATE 40001 (SERIALIZABLE / REPEATABLE READ write conflict)."""
+    return sqlstate_of(exc) == SERIALIZATION_FAILURE
+
+
+def constraint_name(exc: BaseException) -> str:
+    """Return the offending constraint name for a PostgreSQL error, or "" if not applicable."""
+    return constraint_name_of(exc)