sqlpush 0.1.0__py3-none-any.whl

sqlpush/__init__.py

@@ -0,0 +1,78 @@
+from typing import TYPE_CHECKING, Any
+
+from sqlpush.types import (
+    AppliedOperation,
+    CheckResult,
+    ConnectFailed,
+    DestructiveBlocked,
+    MetadataImportError,
+    Plan,
+    PlannedOperation,
+    Report,
+    RiskClass,
+    SqlpushError,
+)
+
+# Static-checker visibility for the PEP 562 lazy exports below: this
+# block never executes at runtime, so the light-import guarantee is
+# untouched (tests/test_annotations.py still guards it).
+if TYPE_CHECKING:
+    from sqlpush.api import (
+        acheck,
+        aensure_schema,
+        aplan,
+        apush,
+        check,
+        ensure_schema,
+        plan,
+        push,
+    )
+
+__version__ = "0.1.0"
+
+# Public API (plan/push/check/ensure_schema + async facade) is exported
+# lazily via PEP 562: an eager import would pull alembic (through
+# sqlpush.core.diff) into every `import sqlpush.*`, breaking the
+# light-import guarantee of the annotations module (see
+# tests/test_annotations.py::test_annotations_module_has_no_heavy_imports).
+_LAZY_API = (
+    "plan",
+    "push",
+    "check",
+    "ensure_schema",
+    "aplan",
+    "apush",
+    "acheck",
+    "aensure_schema",
+)
+
+
+def __getattr__(name: str) -> Any:
+    if name in _LAZY_API:
+        from sqlpush import api
+
+        return getattr(api, name)
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+
+__all__ = [
+    "AppliedOperation",
+    "CheckResult",
+    "ConnectFailed",
+    "DestructiveBlocked",
+    "MetadataImportError",
+    "Plan",
+    "PlannedOperation",
+    "Report",
+    "RiskClass",
+    "SqlpushError",
+    "__version__",
+    "acheck",
+    "aensure_schema",
+    "aplan",
+    "apush",
+    "check",
+    "ensure_schema",
+    "plan",
+    "push",
+]

sqlpush/annotations.py

@@ -0,0 +1,37 @@
+# src/sqlpush/annotations.py
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+HYPERTABLE_KEY = "sqlpush_hypertable"
+
+
+@dataclass(frozen=True)
+class HypertableInfo:
+    time_column: str
+    chunk_time_interval: str | None = None
+
+
+def hypertable(*, time_column: str, chunk_time_interval: str | None = None):
+    """Record hypertable intent on the model's Table (MetaData level).
+
+    Works with SQLModel, Flask-SQLAlchemy and plain declarative, anything
+    whose class already carries a built ``__table__`` when decorated.
+
+    The generated ``create_hypertable`` runs with
+    ``create_default_indexes => false``: timescale's implicit time-column
+    index is invisible to metadata and would drift forever. Declare it
+    yourself (``Index(..., "<time_column>")`` on the model) if you rely
+    on it for time-range scans.
+    """
+
+    def decorator(cls):
+        table = getattr(cls, "__table__", None)
+        if table is None:
+            raise TypeError(f"{cls.__name__} has no __table__; decorate a mapped class")
+        table.info[HYPERTABLE_KEY] = HypertableInfo(
+            time_column=time_column, chunk_time_interval=chunk_time_interval
+        )
+        return cls
+
+    return decorator

sqlpush/api.py

@@ -0,0 +1,172 @@
+# src/sqlpush/api.py
+from __future__ import annotations
+
+import asyncio
+from typing import NoReturn
+
+from sqlalchemy import MetaData
+from sqlalchemy.engine import Engine
+from sqlalchemy.exc import (
+    ArgumentError,
+    OperationalError,
+    SQLAlchemyError,
+)
+from sqlalchemy.ext.asyncio import AsyncEngine
+
+from sqlpush.apply.executor import apply_plan, with_advisory_lock
+from sqlpush.core.diff import DiffEngine
+from sqlpush.directives.timescale import hypertable_operations
+from sqlpush.types import (
+    CheckResult,
+    ConnectFailed,
+    Plan,
+    Report,
+    SqlpushError,
+)
+
+_engine = DiffEngine()
+
+# Connection-phase failures (server unreachable, bad DSN, auth) re-type
+# as ConnectFailed; every other SQLAlchemyError leaving a DB-touching
+# verb becomes a plain SqlpushError. Either way only typed errors
+# escape the API.
+_CONNECT_PHASE = (OperationalError, ArgumentError)
+
+
+def _raise_typed(exc: SQLAlchemyError) -> NoReturn:
+    if isinstance(exc, _CONNECT_PHASE):
+        raise ConnectFailed(f"could not connect to the database: {exc}") from exc
+    raise SqlpushError(f"database error: {exc}") from exc
+
+
+def _build_plan(metadata: MetaData, engine: Engine, schemas, exclude) -> Plan:
+    p = _engine.plan(metadata, engine, schemas=schemas, exclude=exclude)
+    return Plan(operations=p.operations + tuple(hypertable_operations(metadata, engine)))
+
+
+class _PlannerWithDirectives(DiffEngine):
+    """DiffEngine facade whose plans carry directive ops.
+
+    ``with_advisory_lock`` re-plans via ``reverify.plan(metadata, engine,
+    schemas=..., exclude=...)`` once the lock is won; routing that call
+    through the plan builder keeps directive operations (e.g.
+    create_hypertable) on the DEFAULT locked push path.
+    """
+
+    def __init__(self, engine: DiffEngine) -> None:
+        self._engine = engine
+
+    def plan(self, metadata, engine, *, schemas=None, exclude=()) -> Plan:
+        return _build_plan(metadata, engine, schemas, exclude)
+
+
+def plan(metadata, engine, *, schemas=None, exclude=()) -> Plan:
+    try:
+        return _build_plan(metadata, engine, schemas, exclude)
+    except SQLAlchemyError as exc:
+        _raise_typed(exc)
+
+
+def push(
+    metadata,
+    engine,
+    *,
+    safe_only=False,
+    allow_destructive=False,
+    lock=True,
+    lock_timeout=5.0,
+    advisory_wait=30.0,
+    schemas=None,
+    exclude=(),
+) -> Report:
+    try:
+        if lock:
+            return with_advisory_lock(
+                engine,
+                metadata,
+                wait=advisory_wait,
+                timeout=lock_timeout,
+                allow_destructive=allow_destructive,
+                safe_only=safe_only,
+                reverify=_PlannerWithDirectives(_engine),
+                schemas=schemas,
+                exclude=exclude,
+            )
+        p = _build_plan(metadata, engine, schemas, exclude)
+        return apply_plan(
+            engine,
+            p,
+            allow_destructive=allow_destructive,
+            safe_only=safe_only,
+            lock_timeout=lock_timeout,
+        )
+    except SQLAlchemyError as exc:
+        _raise_typed(exc)
+
+
+def check(metadata, engine, *, schemas=None, exclude=()) -> CheckResult:
+    try:
+        p = _build_plan(metadata, engine, schemas, exclude)
+    except SQLAlchemyError as exc:
+        _raise_typed(exc)
+    return CheckResult(
+        clean=not p.drift,
+        drift=p.drift,
+        has_destructive=p.has_destructive,
+    )
+
+
+def _sync_engine_from(target):
+    if isinstance(target, AsyncEngine):
+        from sqlalchemy import create_engine
+        from sqlalchemy.pool import NullPool
+
+        dsn = target.url.render_as_string(hide_password=False)
+        return create_engine(dsn, poolclass=NullPool), True
+    if isinstance(target, str):
+        from sqlalchemy import create_engine
+        from sqlalchemy.pool import NullPool
+
+        return create_engine(target, poolclass=NullPool), True
+    return target, False
+
+
+_ENSURE_MODES = ("push", "check")
+
+
+def ensure_schema(metadata, target, mode="push", **kwargs) -> Report:
+    # validated BEFORE any engine work: a typo'd mode must never fall
+    # through to the writing branch
+    if mode not in _ENSURE_MODES:
+        raise SqlpushError(f"invalid mode {mode!r}: expected 'push' or 'check'")
+    try:
+        engine, dispose = _sync_engine_from(target)
+        try:
+            if mode == "check":
+                result = check(metadata, engine, **kwargs)
+                if result.drift:
+                    raise SqlpushError("schema drift detected (ensure_schema mode='check')")
+                return Report()
+            return push(metadata, engine, **kwargs)
+        finally:
+            if dispose:
+                engine.dispose()
+    except SQLAlchemyError as exc:
+        # includes ArgumentError from a bad string DSN in _sync_engine_from
+        _raise_typed(exc)
+
+
+async def aplan(metadata, engine, **kw):
+    return await asyncio.to_thread(plan, metadata, engine, **kw)
+
+
+async def apush(metadata, engine, **kw):
+    return await asyncio.to_thread(push, metadata, engine, **kw)
+
+
+async def acheck(metadata, engine, **kw):
+    return await asyncio.to_thread(check, metadata, engine, **kw)
+
+
+async def aensure_schema(metadata, target, **kw):
+    return await asyncio.to_thread(ensure_schema, metadata, target, **kw)

sqlpush/apply/executor.py

@@ -0,0 +1,199 @@
+# src/sqlpush/apply/executor.py
+from __future__ import annotations
+
+import contextlib
+import time
+from typing import TYPE_CHECKING
+
+from sqlalchemy import text
+from sqlalchemy.engine import Engine
+
+from sqlpush.types import (
+    AppliedOperation,
+    Plan,
+    PlannedOperation,
+    Report,
+    RiskClass,
+    SqlpushError,
+)
+
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+
+    from sqlalchemy import MetaData
+
+    from sqlpush.core.diff import DiffEngine
+
+
+def _is_concurrent(op: PlannedOperation) -> bool:
+    return "CONCURRENTLY" in op.sql.upper()
+
+
+def apply_plan(
+    engine: Engine,
+    plan: Plan,
+    *,
+    allow_destructive: bool = False,
+    safe_only: bool = False,
+    lock_timeout: float = 5.0,
+) -> Report:
+    start = time.monotonic()
+
+    blocked = tuple(op for op in plan.operations if op.risk is RiskClass.DESTRUCTIVE)
+    if blocked and not allow_destructive:
+        # Destructive gate: nothing executes at all. When safe_only is also
+        # active, risky ops are still recorded as policy-skipped; destructive
+        # ones are already in `blocked` and are not double-listed.
+        skipped = (
+            tuple(op for op in plan.operations if op.risk is RiskClass.RISKY) if safe_only else ()
+        )
+        return Report(
+            applied=(),
+            blocked=blocked,
+            skipped=skipped,
+            partial_failure=False,
+            duration=time.monotonic() - start,
+        )
+
+    if safe_only:
+        runnable = [op for op in plan.operations if op.risk is RiskClass.SAFE]
+        skipped = tuple(op for op in plan.operations if op.risk is not RiskClass.SAFE)
+    else:
+        runnable = list(plan.operations)
+        skipped = ()
+
+    applied: list[AppliedOperation] = []
+    partial_failure = False
+
+    # --- CONCURRENTLY segment: autocommit, one op per statement ----------
+    concurrent = [op for op in runnable if _is_concurrent(op)]
+    if concurrent:
+        with engine.connect().execution_options(isolation_level="AUTOCOMMIT") as conn:
+            for op in concurrent:
+                try:
+                    conn.execute(text(op.sql))
+                    applied.append(AppliedOperation(op.type, "applied"))
+                except Exception:  # noqa: BLE001  # concurrent ops record any DB failure and continue
+                    applied.append(AppliedOperation(op.type, "failed"))
+                    partial_failure = True
+
+    # --- transactional segment: atomic ----------------------------------
+    plain = [op for op in runnable if not _is_concurrent(op)]
+    if plain:
+        try:
+            with engine.begin() as conn:
+                # NOTE: PostgreSQL does not accept bind parameters for SET
+                # (utility statement), so the int is inlined; lock_timeout is
+                # a typed float parameter, not user input.
+                conn.execute(text(f"SET LOCAL lock_timeout = {int(lock_timeout * 1000)}"))
+                for op in plain:
+                    conn.execute(text(op.sql))
+            applied.extend(AppliedOperation(op.type, "applied") for op in plain)
+        except Exception as exc:
+            msg = f"transactional segment failed, rolled back: {exc}"
+            if applied:
+                msg += f"; concurrent segment had already applied: {[a.type for a in applied]}"
+            raise SqlpushError(msg) from exc
+
+    return Report(
+        applied=tuple(applied),
+        blocked=(),
+        skipped=skipped,
+        partial_failure=partial_failure,
+        duration=time.monotonic() - start,
+    )
+
+
+# --- advisory lock: winner/loser semantics --------------------------------
+#
+# v0.1 design notes:
+# - The lock is session-scoped and taken/released on `conn`, while
+#   `apply_plan` opens its own connections. Acceptable because the
+#   advisory lock's job is worker coordination (one pusher per database
+#   at a time), not txn participation.
+# - Session advisory locks survive ROLLBACK: the winner rolls back the
+#   txn opened by the key/probe queries right after acquiring, so the
+#   session never idles in transaction (an
+#   idle_in_transaction_session_timeout would kill the winner mid-push
+#   and silently release the lock).
+
+
+def fnv1a_32(data: bytes) -> int:
+    h = 2166136261
+    for byte in data:
+        h ^= byte
+        h = (h * 16777619) & 0xFFFFFFFF
+    return h
+
+
+def advisory_key(conn) -> int:
+    # Deterministic across DSN spellings: the key derives from the
+    # database's oid, not from the connection string.
+    oid = conn.execute(
+        text("SELECT oid FROM pg_database WHERE datname = current_database()")
+    ).scalar()
+    return fnv1a_32(b"sqlpush") ^ oid
+
+
+def with_advisory_lock(
+    engine: Engine,
+    metadata: MetaData,
+    *,
+    wait: float = 30.0,
+    timeout: float = 5.0,
+    allow_destructive: bool = False,
+    safe_only: bool = False,
+    reverify: DiffEngine | None = None,
+    schemas: Sequence[str] | None = None,
+    exclude: Sequence[str] = (),
+) -> Report:
+    """Winner migrates; losers block (bounded), then re-verify.
+
+    Losers poll ``pg_try_advisory_lock`` every 0.5 s against a
+    ``time.monotonic()`` deadline; once the lock is acquired the winner
+    path re-plans (covering the case where the previous winner died
+    mid-push). Raises :class:`SqlpushError` if the wait budget is
+    exhausted.
+    """
+    if reverify is None:
+        raise SqlpushError(
+            "reverify is required: pass an object exposing a "
+            "DiffEngine-compatible .plan(metadata, engine, ...)"
+        )
+    if wait < 0:
+        raise SqlpushError(f"wait must be >= 0, got {wait}")
+    if timeout < 0:
+        raise SqlpushError(f"timeout (lock_timeout) must be >= 0, got {timeout}")
+    deadline = time.monotonic() + wait
+    with engine.connect() as conn:
+        key = advisory_key(conn)
+        locked = conn.execute(text("SELECT pg_try_advisory_lock(:k)"), {"k": key}).scalar()
+        while not locked and time.monotonic() < deadline:
+            time.sleep(0.5)
+            locked = conn.execute(text("SELECT pg_try_advisory_lock(:k)"), {"k": key}).scalar()
+        if not locked:
+            raise SqlpushError(f"another sqlpush worker holds the advisory lock after {wait}s")
+        # close the txn opened by the key/probe queries: session
+        # advisory locks survive ROLLBACK, and an idle-in-transaction
+        # winner is exposed to idle_in_transaction_session_timeout
+        # (killed mid-push = lock silently released)
+        conn.rollback()
+        try:
+            plan = reverify.plan(metadata, engine, schemas=schemas, exclude=exclude)
+            if not plan.drift:
+                return Report()
+            return apply_plan(
+                engine,
+                plan,
+                allow_destructive=allow_destructive,
+                safe_only=safe_only,
+                lock_timeout=timeout,
+            )
+        finally:
+            # best-effort unlock: a secondary failure here (e.g. the
+            # txn is already aborted) must never mask the original
+            # exception out of the winner path
+            if not conn.closed:
+                with contextlib.suppress(Exception):
+                    conn.rollback()
+                    conn.execute(text("SELECT pg_advisory_unlock(:k)"), {"k": key})

sqlpush/cli.py

@@ -0,0 +1,218 @@
+# src/sqlpush/cli.py
+"""sqlpush: diff/push/check PostgreSQL schema drift from SQLAlchemy models.
+
+Exit codes: diff always 0; check 0 clean / 2 drift / 3 destructive drift;
+push 0 applied / 1 destructive blocked / 2 error (incl. partial failure).
+"""
+
+from __future__ import annotations
+
+import importlib
+import json
+import os
+import sys
+from typing import Annotated
+
+import typer
+from sqlalchemy import create_engine
+from sqlalchemy.exc import ArgumentError
+from sqlalchemy.pool import NullPool
+
+from sqlpush import api
+from sqlpush.core.render import render
+from sqlpush.types import (
+    CheckResult,
+    MetadataImportError,
+    RiskClass,
+    SqlpushError,
+)
+
+app = typer.Typer(add_completion=False, help=__doc__ or "")
+
+# Annotated + None default: keeps ruff B008/B006 quiet (no call in the
+# default, no mutable default) while staying the idiomatic typer style.
+SchemaOpt = Annotated[list[str] | None, typer.Option("--schema")]
+ExcludeOpt = Annotated[list[str] | None, typer.Option("--exclude")]
+
+
+def _load_metadata(spec: str):
+    module, _, attr = spec.partition(":")
+    try:
+        obj = importlib.import_module(module)
+        for part in attr.split("."):
+            obj = getattr(obj, part)
+        return obj
+    except (ImportError, AttributeError) as exc:
+        raise MetadataImportError(f"cannot import {spec!r}: {exc}") from exc
+
+
+def _engine(dsn: str | None):
+    dsn = dsn or os.environ.get("DATABASE_URL")
+    if not dsn:
+        typer.secho("no --dsn and no DATABASE_URL set", fg="red", err=True)
+        raise typer.Exit(code=1)
+    try:
+        return create_engine(dsn, poolclass=NullPool)
+    except ArgumentError as exc:
+        # malformed DSN is a configuration error, not a verb error: exit 1
+        typer.secho(f"invalid DSN: {exc}", fg="red", err=True)
+        raise typer.Exit(code=1) from exc
+
+
+def _emit_json(plan) -> None:
+    typer.echo(json.dumps(plan.to_json_dict(), indent=2))
+
+
+def _risk_summary(p) -> None:
+    # --verbose adds an op-count line per PRESENT risk class (absent
+    # classes stay silent, mirroring render's section omission)
+    for cls in RiskClass:
+        n = sum(1 for op in p.operations if op.risk is cls)
+        if n:
+            typer.echo(f"{cls.value}: {n} operation(s)")
+
+
+@app.command()
+def diff(
+    metadata_spec: str = typer.Argument(..., help="module:metadata"),
+    dsn: str | None = typer.Option(None),
+    json_output: bool = typer.Option(False, "--json"),
+    verbose: bool = typer.Option(False, "--verbose"),
+    quiet: bool = typer.Option(False, "--quiet"),
+    schema: SchemaOpt = None,
+    exclude: ExcludeOpt = None,
+):
+    md = _load_metadata(metadata_spec)
+    engine = _engine(dsn)
+    try:
+        p = api.plan(md, engine, schemas=schema, exclude=exclude or ())
+        if json_output:
+            _emit_json(p)
+            raise typer.Exit(code=0)
+        if verbose:
+            _risk_summary(p)
+        typer.echo(render(p) or "-- schema in sync --")
+    finally:
+        engine.dispose()
+    raise typer.Exit(code=0)
+
+
+@app.command()
+def check(
+    metadata_spec: str = typer.Argument(...),
+    dsn: str | None = typer.Option(None),
+    json_output: bool = typer.Option(False, "--json"),
+    verbose: bool = typer.Option(False, "--verbose"),
+    quiet: bool = typer.Option(False, "--quiet"),
+    schema: SchemaOpt = None,
+    exclude: ExcludeOpt = None,
+):
+    md = _load_metadata(metadata_spec)
+    engine = _engine(dsn)
+    try:
+        # plan ONCE and derive everything from that single object: a
+        # second plan could observe a different DB state than the one
+        # the exit code was derived from (TOCTOU)
+        p = api.plan(md, engine, schemas=schema, exclude=exclude or ())
+        result = CheckResult(
+            clean=not p.drift,
+            drift=p.drift,
+            has_destructive=p.has_destructive,
+        )
+        if json_output:
+            _emit_json(p)
+        elif verbose:
+            ndest = sum(1 for op in p.operations if op.risk is RiskClass.DESTRUCTIVE)
+            typer.echo(f"drift: {len(p.operations)} operation(s), {ndest} destructive")
+    finally:
+        engine.dispose()
+    if result.clean:
+        raise typer.Exit(code=0)
+    raise typer.Exit(code=3 if result.has_destructive else 2)
+
+
+@app.command()
+def push(
+    metadata_spec: str = typer.Argument(...),
+    dsn: str | None = typer.Option(None),
+    json_output: bool = typer.Option(False, "--json"),
+    allow_destructive: bool = typer.Option(False, "--allow-destructive"),
+    safe_only: bool = typer.Option(False, "--safe-only"),
+    no_lock: bool = typer.Option(False, "--no-lock"),
+    lock_timeout: float = typer.Option(5.0, "--lock-timeout"),
+    advisory_wait: float = typer.Option(30.0, "--advisory-wait"),
+    verbose: bool = typer.Option(False, "--verbose"),
+    quiet: bool = typer.Option(False, "--quiet"),
+    schema: SchemaOpt = None,
+    exclude: ExcludeOpt = None,
+):
+    md = _load_metadata(metadata_spec)
+    engine = _engine(dsn)
+    try:
+        try:
+            report = api.push(
+                md,
+                engine,
+                safe_only=safe_only,
+                allow_destructive=allow_destructive,
+                lock=not no_lock,
+                lock_timeout=lock_timeout,
+                advisory_wait=advisory_wait,
+                schemas=schema,
+                exclude=exclude or (),
+            )
+        except SqlpushError as exc:
+            # binding exit-code table: 1 is exclusively "destructive
+            # blocked"; every other push error (lock timeout, connect
+            # failure, ...) is exit 2
+            typer.secho(str(exc), fg="red", err=True)
+            raise typer.Exit(code=2) from exc
+        if json_output:
+            typer.echo(
+                json.dumps(
+                    {
+                        "version": 1,
+                        "applied": [a.__dict__ for a in report.applied],
+                        "blocked": [b.__dict__ for b in report.blocked],
+                        "skipped": [s.__dict__ for s in report.skipped],
+                        "partial_failure": report.partial_failure,
+                    },
+                    indent=2,
+                    default=str,
+                )
+            )
+        elif verbose:
+            for a in report.applied:
+                typer.echo(f"{a.status}: {a.type}")
+            if report.skipped:
+                typer.echo(f"skipped: {len(report.skipped)} operation(s)")
+            typer.echo(f"duration: {report.duration:.3f}s")
+    finally:
+        engine.dispose()
+    if report.blocked:
+        if not quiet:
+            typer.secho(
+                f"{len(report.blocked)} destructive operation(s) blocked; "
+                "re-run with --allow-destructive",
+                fg="yellow",
+                err=True,
+            )
+        raise typer.Exit(code=1)
+    if report.skipped and not json_output and not quiet:
+        # policy-skipped (--safe-only) is informational, never an error
+        typer.secho(
+            f"{len(report.skipped)} operation(s) skipped by --safe-only",
+            fg="blue",
+            err=True,
+        )
+    if report.partial_failure:
+        raise typer.Exit(code=2)
+    raise typer.Exit(code=0)
+
+
+def main() -> None:  # [project.scripts] entry point
+    try:
+        app()
+    except (MetadataImportError, SqlpushError) as exc:
+        typer.secho(str(exc), fg="red", err=True)
+        sys.exit(1)