SPB Git

spb/findata

Public
6.0 KB · 114 lines python
1"""Durable job queue on PostgreSQL (SELECT … FOR UPDATE SKIP LOCKED). Survives reboots; retries with2exponential backoff + jitter; dead-letter after max_attempts; checkpoints keyed by scope."""3from __future__ import annotations45import json6import random7import socket8from datetime import UTC, datetime, timedelta9from typing import Any1011from findata_core.config import get_settings12from findata_core.db import connection, transaction13from findata_core.redact import redact1415WORKER_ID = f"{socket.gethostname()}:{random.randint(1000, 9999)}"161718def enqueue(kind: str, dataset_id: str | None = None, params: dict[str, Any] | None = None, priority: int = 50,19            dedupe_key: str | None = None, max_attempts: int = 5, created_by: str = "system",20            run_at: datetime | None = None, rerun_succeeded: bool = True) -> str | None:21    """Insert a job; returns id, or None if a job with the same dedupe_key is already queued/running.22    With rerun_succeeded=False (scheduler), a key whose job already succeeded is not re-enqueued either —23    this is what makes daily keys like `quality:2026-11-02` fire once per day."""24    with transaction() as conn:25        if dedupe_key:26            r = conn.execute("SELECT id, status FROM ingest_jobs WHERE dedupe_key=%s", (dedupe_key,)).fetchone()27            if r:28                if r["status"] in ("queued", "running"):29                    return None30                if r["status"] == "succeeded" and not rerun_succeeded:31                    return None32                conn.execute("UPDATE ingest_jobs SET dedupe_key=NULL WHERE id=%s", (r["id"],))33        r = conn.execute(34            "INSERT INTO ingest_jobs(kind, dataset_id, params, priority, dedupe_key, max_attempts, created_by, next_run_at) "35            "VALUES (%s,%s,%s,%s,%s,%s,%s,%s) RETURNING id",36            (kind, dataset_id, json.dumps(params or {}), priority, dedupe_key, max_attempts, created_by,37             run_at or datetime.now(UTC))).fetchone()38        return str(r["id"])394041def claim_next(kinds: list[str] | None = None) -> dict[str, Any] | None:42    with transaction() as conn:43        sql = ("SELECT * FROM ingest_jobs WHERE status='queued' AND next_run_at <= now() "44               + ("AND kind = ANY(%s) " if kinds else "") +45               "ORDER BY priority, next_run_at LIMIT 1 FOR UPDATE SKIP LOCKED")46        job = conn.execute(sql, (kinds,) if kinds else None).fetchone()47        if not job:48            return None49        conn.execute("UPDATE ingest_jobs SET status='running', attempts=attempts+1, locked_by=%s, locked_at=now(), "50                     "started_at=COALESCE(started_at, now()) WHERE id=%s", (WORKER_ID, job["id"]))51        run = conn.execute("INSERT INTO ingest_runs(job_id, worker, code_version) VALUES (%s,%s,%s) RETURNING id",52                           (job["id"], WORKER_ID, _code_version())).fetchone()53        job = dict(job)54        job["attempts"] += 155        job["run_id"] = str(run["id"])56        return job575859def finish(job: dict[str, Any], ok: bool, result: dict[str, Any] | None = None, error: str | None = None) -> None:60    secrets = get_settings().secret_values()61    with transaction() as conn:62        if ok:63            conn.execute("UPDATE ingest_jobs SET status='succeeded', finished_at=now(), locked_by=NULL, result=%s, last_error=NULL WHERE id=%s",64                         (json.dumps(result or {}, default=str), job["id"]))65            conn.execute("UPDATE ingest_runs SET status='succeeded', finished_at=now(), stats=%s WHERE id=%s",66                         (json.dumps(result or {}, default=str), job["run_id"]))67            return68        msg = redact(error or "unknown error", secrets)[:4000]69        dead = job["attempts"] >= job["max_attempts"]70        if dead:71            conn.execute("UPDATE ingest_jobs SET status='dead', finished_at=now(), locked_by=NULL, last_error=%s WHERE id=%s", (msg, job["id"]))72        else:73            delay = min(3600, 30 * 2 ** job["attempts"]) + random.uniform(0, 30)74            conn.execute("UPDATE ingest_jobs SET status='queued', locked_by=NULL, last_error=%s, next_run_at=%s WHERE id=%s",75                         (msg, datetime.now(UTC) + timedelta(seconds=delay), job["id"]))76        conn.execute("UPDATE ingest_runs SET status=%s, finished_at=now(), stats=%s WHERE id=%s",77                     ("dead" if dead else "failed", json.dumps({"error": msg}), job["run_id"]))787980def progress(job_id: str, data: dict[str, Any]) -> None:81    """Store progress and refresh the heartbeat (locked_at) used by requeue_stale()."""82    with transaction() as conn:83        conn.execute("UPDATE ingest_jobs SET progress=%s, locked_at=now() WHERE id=%s", (json.dumps(data, default=str), job_id))848586def set_checkpoint(scope: str, key: str, value: dict[str, Any]) -> None:87    with transaction() as conn:88        conn.execute("INSERT INTO checkpoints(scope, key, value) VALUES (%s,%s,%s) ON CONFLICT (scope, key) DO UPDATE SET value=EXCLUDED.value, updated_at=now()",89                     (scope, key, json.dumps(value, default=str)))909192def get_checkpoint(scope: str, key: str) -> dict[str, Any] | None:93    with connection() as conn:94        r = conn.execute("SELECT value FROM checkpoints WHERE scope=%s AND key=%s", (scope, key)).fetchone()95        return r["value"] if r else None969798def requeue_stale(max_age_minutes: int = 90) -> int:99    """Jobs 'running' without a heartbeat (progress) for `max_age_minutes` are re-queued — the worker is gone100    (idempotent tasks + checkpoints make this safe). Long jobs stay alive by reporting progress."""101    with transaction() as conn:102        r = conn.execute("UPDATE ingest_jobs SET status='queued', locked_by=NULL, next_run_at=now() WHERE status='running' "103                         "AND locked_at < now() - make_interval(mins => %s) RETURNING id", (max_age_minutes,)).fetchall()104        return len(r)105106107def _code_version() -> str:108    import subprocess109    try:110        return subprocess.run(["git", "rev-parse", "--short", "HEAD"], capture_output=True, text=True, timeout=5,111                              cwd=str(get_settings().root)).stdout.strip() or "unknown"112    except Exception:113        return "unknown"114
Files · main
Menu

Theme