feat: add authenticated engine processing boundary

site/db/extraction_jobs.py

@@ -0,0 +1,79 @@
+"""Persistent extraction progress and per-batch worker ownership."""
+import json
+import time
+import uuid
+
+from db.connection import execute, query
+
+LEASE_SECONDS = 300
+
+
+def claim(batch_id, total):
+    """Claim a batch atomically, allowing retry after failure or expired lease."""
+    owner = str(uuid.uuid4())
+    now = time.time()
+    changed = execute(
+        """INSERT INTO extraction_jobs (batch_id, owner, state, total, updated_at)
+           VALUES (%s, %s, 'running', %s, %s)
+           ON CONFLICT(batch_id) DO UPDATE SET owner=excluded.owner, state='running',
+             total=excluded.total, done=0, failed=0, classes_json='{}',
+             result_json=NULL, error=NULL, updated_at=excluded.updated_at
+           WHERE extraction_jobs.state != 'running' OR extraction_jobs.updated_at < %s""",
+        (batch_id, owner, total, now, now - LEASE_SECONDS),
+    )
+    return owner if changed else None
+
+
+def heartbeat(batch_id, owner):
+    return execute(
+        "UPDATE extraction_jobs SET updated_at=%s WHERE batch_id=%s AND owner=%s AND state='running'",
+        (time.time(), batch_id, owner),
+    )
+
+
+def progress(batch_id, owner, completed, result):
+    rows = query(
+        "SELECT failed, classes_json FROM extraction_jobs WHERE batch_id=%s AND owner=%s AND state='running'",
+        (batch_id, owner),
+    )
+    if not rows:
+        return
+    failed = rows[0]["failed"]
+    classes = json.loads(rows[0]["classes_json"])
+    if not result.get("ok"):
+        failed += 1
+    else:
+        doc_class = (result.get("extraction") or {}).get("doc_class", "unknown")
+        classes[doc_class] = classes.get(doc_class, 0) + 1
+    execute(
+        """UPDATE extraction_jobs SET done=%s, failed=%s, classes_json=%s, updated_at=%s
+           WHERE batch_id=%s AND owner=%s AND state='running'""",
+        (completed, failed, json.dumps(classes), time.time(), batch_id, owner),
+    )
+
+
+def finish(batch_id, owner, result=None, error=None):
+    state = "done" if result and result.get("ok") and not error else "failed"
+    return execute(
+        """UPDATE extraction_jobs SET state=%s, result_json=%s, error=%s, updated_at=%s
+           WHERE batch_id=%s AND owner=%s AND state='running'""",
+        (state, json.dumps(result, ensure_ascii=False) if result is not None else None,
+         error or ("extraction failed" if state == "failed" else None), time.time(), batch_id, owner),
+    )
+
+
+def get(batch_id):
+    """Mark an abandoned lease interrupted; callers may explicitly retry it."""
+    execute(
+        """UPDATE extraction_jobs SET state='interrupted', error='worker stopped; retry the batch'
+           WHERE batch_id=%s AND state='running' AND updated_at < %s""",
+        (batch_id, time.time() - LEASE_SECONDS),
+    )
+    rows = query("SELECT * FROM extraction_jobs WHERE batch_id=%s", (batch_id,))
+    if not rows:
+        return None
+    job = rows[0]
+    job["classes"] = json.loads(job.pop("classes_json"))
+    result_json = job.pop("result_json")
+    job["result"] = json.loads(result_json) if result_json else None
+    return job