"""Postgres connection + plain-SQL migrations + job queue helpers. No ORM. The schema in migrations/ is the contract (SPEC §3, §6). Job queue semantics: - enqueue(type, payload, run_after) inserts a 'queued' row. - claim_next_job() atomically claims one due job (FOR UPDATE SKIP LOCKED) and marks it 'running'. - tick() (in cli.py) dispatches claimed jobs to registered handlers and marks them done/failed. Failures are loud: status='failed' + error text (SPEC §9). """ from __future__ import annotations import json import re from contextlib import contextmanager from pathlib import Path from typing import Any, Iterator import psycopg from psycopg.rows import dict_row from psycopg.types.json import Json from .config import REPO_ROOT, get_settings MIGRATIONS_DIR = REPO_ROOT / "migrations" _MIGRATION_FILE_RE = re.compile(r"^(\d+)_.+\.sql$") def schema() -> str: """The Postgres schema v2 lives in (DB_SCHEMA). Validated to a plain identifier so it is safe to interpolate into search_path / DDL.""" s = (get_settings().db_schema or "public").strip() if not re.fullmatch(r"[A-Za-z_][A-Za-z0-9_]*", s): raise RuntimeError(f"invalid DB_SCHEMA {s!r} — use a simple identifier (letters/digits/_)") return s @contextmanager def connect(autocommit: bool = False) -> Iterator[psycopg.Connection]: conn = psycopg.connect( get_settings().database_url, row_factory=dict_row, autocommit=autocommit, options=f"-c search_path={schema()},public", ) try: yield conn if not autocommit: conn.commit() except Exception: if not autocommit: conn.rollback() raise finally: conn.close() # ── Migrations ──────────────────────────────────────────────────────────────── def migrate() -> list[str]: """Apply pending numbered migrations in order. Returns the filenames applied.""" files = sorted(p for p in MIGRATIONS_DIR.glob("*.sql") if _MIGRATION_FILE_RE.match(p.name)) applied: list[str] = [] with connect() as conn: conn.execute(f'create schema if not exists "{schema()}"') conn.execute( """ create table if not exists schema_migrations ( filename text primary key, applied_at timestamptz not null default now() ) """ ) done = { r["filename"] for r in conn.execute("select filename from schema_migrations").fetchall() } for path in files: if path.name in done: continue conn.execute(path.read_text()) # type: ignore[arg-type] conn.execute( "insert into schema_migrations (filename) values (%s)", (path.name,) ) applied.append(path.name) return applied # ── Job queue ───────────────────────────────────────────────────────────────── def enqueue(job_type: str, payload: dict[str, Any] | None = None, run_after: str | None = None) -> str: """Insert a queued job. run_after is an optional SQL-parseable timestamp.""" with connect() as conn: row = conn.execute( """ insert into jobs (type, payload, run_after) values (%s, %s, coalesce(%s::timestamptz, now())) returning id """, (job_type, Json(payload or {}), run_after), ).fetchone() assert row is not None return str(row["id"]) def claim_next_job(conn: psycopg.Connection) -> dict[str, Any] | None: """Atomically claim one due queued job and mark it running. None if no work.""" row = conn.execute( """ update jobs set status = 'running', attempts = attempts + 1 where id = ( select id from jobs where status = 'queued' and run_after <= now() order by run_after, created_at for update skip locked limit 1 ) returning id, type, payload, attempts """ ).fetchone() return dict(row) if row else None def complete_job(conn: psycopg.Connection, job_id: str) -> None: conn.execute("update jobs set status = 'done', error = null where id = %s", (job_id,)) def fail_job(conn: psycopg.Connection, job_id: str, error: str) -> None: conn.execute("update jobs set status = 'failed', error = %s where id = %s", (error, job_id)) def job_status(job_id: str) -> dict[str, Any] | None: with connect() as conn: row = conn.execute( "select id, type, status, attempts, error from jobs where id = %s", (job_id,) ).fetchone() return dict(row) if row else None # ── Small helpers used across modules ──────────────────────────────────────── def fetch_one(sql: str, params: tuple | None = None) -> dict[str, Any] | None: with connect() as conn: row = conn.execute(sql, params).fetchone() return dict(row) if row else None def fetch_all(sql: str, params: tuple | None = None) -> list[dict[str, Any]]: with connect() as conn: return [dict(r) for r in conn.execute(sql, params).fetchall()] def execute(sql: str, params: tuple | None = None) -> None: with connect() as conn: conn.execute(sql, params) def to_json(value: Any) -> str: return json.dumps(value, ensure_ascii=False)