worker.py

#!/usr/bin/env python3
"""Process queued jobs (vision / extract / embed / pipeline)."""

from __future__ import annotations

import argparse
import sys
import time
from pathlib import Path

from dotenv import load_dotenv

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))
load_dotenv(ROOT / ".env")

from app.db import apply_migrations, repository as repo  # noqa: E402
from app.services.pipeline import enqueue_job, process_job  # noqa: E402


def run_once(job_type: str | None) -> int:
    apply_migrations()
    if job_type:
        tenant = repo.get_tenant_by_slug("demo") or repo.ensure_tenant("demo", "Demo tenant")
        job, replayed = enqueue_job(
            tenant_id=tenant["id"],
            job_type=job_type,
            idempotency_key=f"cli-{job_type}",
        )
        print(f"job id={job['id']} type={job['type']} replayed={replayed} status={job['status']}")
        if job["status"] in {"queued", "failed", "running"}:
            repo.update_job(job["id"], status="running", bump_attempts=True)
            fresh = repo.get_job(job["id"], tenant["id"])
            assert fresh is not None
            process_job(fresh)
        return 0

    claimed = repo.claim_next_job()
    if claimed is None:
        print("no queued jobs")
        return 0
    print(f"claimed job {claimed['id']} type={claimed['type']}")
    process_job(claimed)
    return 0


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--once", action="store_true", help="Process one queued job and exit")
    parser.add_argument(
        "--pipeline",
        action="store_true",
        help="Enqueue and run vision+extract+embed once",
    )
    parser.add_argument("--poll", type=int, default=0, help="Loop every N seconds (0 = no loop)")
    args = parser.parse_args()

    if args.pipeline:
        return run_once("pipeline")
    if args.once or args.poll == 0:
        return run_once(None)
    while True:
        run_once(None)
        time.sleep(args.poll)


if __name__ == "__main__":
    raise SystemExit(main())