pipeline/db.py

"""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)