PacketStation/server/packetstation/api/jobs.py

"""The worker's jobs: schedule, heartbeat, recent runs, and "run now".

Admin (password session) or Clerical (clerical or clerical_supervisor). "Run now" inserts a job_requests row and
NOTIFYs `ps_job_request` in the same transaction; the worker claims every pending request in one run.
`run_out`, `worker_status` and `request_run` are shared with api/feed.py.
"""

from __future__ import annotations

import json
import re
from typing import Any

from fastapi import APIRouter, Depends, Query
from psycopg.rows import dict_row
from pydantic import BaseModel, Field

from .. import audit
from ..worker.heartbeat import STALE_SECONDS
from .deps import Db, Principal, current_principal, get_db, require_password_session
from .errors import ApiError
from .packets import WRITE

router = APIRouter(prefix="/api/jobs", tags=["jobs"])

CLERICAL = frozenset({"clerical", "clerical_supervisor"})
JOB_CHANNEL = "ps_job_request"
_JOB_NAME = re.compile(r"^[a-z0-9][a-z0-9-]{0,39}$")


def jobs_principal(p: Principal = Depends(current_principal)) -> Principal:  # noqa: B008
    """Clerical in any session, or an admin with a password session."""
    if CLERICAL & set(p.roles):
        return p
    if p.is_admin:
        return require_password_session(p)
    raise ApiError(403, "forbidden", "Only Clerical or an admin can see and run the worker's jobs.")


class RunBody(BaseModel):
    note: str | None = Field(None, max_length=500)


def run_out(r: dict[str, Any] | None) -> dict[str, Any] | None:
    """One job_runs row as JSON."""
    if r is None:
        return None
    return {
        "id": r["id"],
        "job": r["job"],
        "trigger": r["trigger"],
        "requested_by": [str(u) for u in (r.get("requested_by") or [])],
        "started_at": r["started_at"],
        "finished_at": r["finished_at"],
        "status": r["status"],
        "counts": r.get("counts") or {},
        "errors": r.get("errors") or [],
        "build_version": r.get("build_version"),
        "visual_target": r.get("visual_target"),
        "host": r.get("host"),
    }


def recent_runs(conn: Any, job: str, limit: int = 10) -> list[dict[str, Any]]:
    with conn.cursor(row_factory=dict_row) as cur:
        rows = cur.execute("SELECT * FROM job_runs WHERE job = %s ORDER BY id DESC LIMIT %s", (job, limit)).fetchall()
    return [o for o in (run_out(r) for r in rows) if o is not None]


def worker_status(conn: Any) -> dict[str, Any]:
    """{at, build, host, running, stale}: stale when the worker has not beaten for STALE_SECONDS or never did."""
    with conn.cursor(row_factory=dict_row) as cur:
        r = cur.execute(
            "SELECT at, build, host, running, now() - at > make_interval(secs => %s) AS stale"
            " FROM worker_heartbeat WHERE id = 1",
            (STALE_SECONDS,),
        ).fetchone()
    if r is None:
        return {"at": None, "build": None, "host": None, "running": [], "stale": True}
    return {
        "at": r["at"],
        "build": r["build"],
        "host": r["host"],
        "running": list(r["running"] or []),
        "stale": bool(r["stale"]),
    }


def known_job(conn: Any, job: str) -> bool:
    if not _JOB_NAME.match(job):
        return False
    return conn.execute("SELECT 1 FROM job_schedule WHERE job = %s", (job,)).fetchone() is not None


def request_run(
    conn: Any, p: Principal, job: str, *, note: str | None = None, action: str = "job.run_requested"
) -> dict[str, Any]:
    """Insert a request and NOTIFY the worker in one transaction. Returns {requested, pending}."""
    with conn.transaction():
        conn.execute(
            "INSERT INTO job_requests (job, requested_by, note) VALUES (%s, %s, %s)",
            (job, p.user_id, (note or "").strip() or None),
        )
        conn.execute("SELECT ps_notify(%s, %s)", (JOB_CHANNEL, json.dumps({"job": job}, separators=(",", ":"))))
        pending = conn.execute(
            "SELECT count(*) FROM job_requests WHERE job = %s AND claimed_at IS NULL", (job,)
        ).fetchone()[0]
        audit.append(
            conn,
            action=action,
            user_id=p.user_id,
            session_id=p.session_id,
            role_snapshot=p.role_snapshot,
            subject_type="job",
            subject_id=job,
            detail={"job": job, "pending": int(pending)},
            ip=p.ip,
        )
    return {"requested": True, "pending": int(pending)}


@router.get("")
def list_jobs(p: Principal = Depends(jobs_principal), d: Db = Depends(get_db)) -> dict[str, Any]:  # noqa: B008
    conn = d.conn
    with conn.cursor(row_factory=dict_row) as cur:
        rows = cur.execute(
            "SELECT s.job, s.next_run_at, s.interval_seconds, s.paused, s.updated_at,"
            " (SELECT count(*) FROM job_requests q WHERE q.job = s.job AND q.claimed_at IS NULL) AS pending,"
            " CASE WHEN r.id IS NULL THEN NULL ELSE to_jsonb(r) END AS last_run"
            " FROM job_schedule s"
            " LEFT JOIN LATERAL (SELECT * FROM job_runs j WHERE j.job = s.job ORDER BY j.id DESC LIMIT 1) r ON true"
            " ORDER BY s.job"
        ).fetchall()
    jobs = [
        {
            "job": r["job"],
            "next_run_at": r["next_run_at"],
            "interval_seconds": r["interval_seconds"],
            "paused": r["paused"],
            "updated_at": r["updated_at"],
            "pending": int(r["pending"]),
            "last_run": run_out(r["last_run"]),
        }
        for r in rows
    ]
    return {"jobs": jobs, "heartbeat": worker_status(conn)}


@router.get("/{job}/runs")
def job_runs(
    job: str,
    p: Principal = Depends(jobs_principal),  # noqa: B008
    d: Db = Depends(get_db),  # noqa: B008
    limit: int = Query(10, ge=1, le=100),
) -> list[dict[str, Any]]:
    if not known_job(d.conn, job):
        raise ApiError(404, "not_found", "No such job.")
    return recent_runs(d.conn, job, limit)


@router.post("/{job}/run", dependencies=WRITE)
def run_now(
    job: str,
    body: RunBody | None = None,
    p: Principal = Depends(jobs_principal),  # noqa: B008
    d: Db = Depends(get_db),  # noqa: B008
) -> dict[str, Any]:
    if not known_job(d.conn, job):
        raise ApiError(404, "not_found", "No such job.")
    return request_run(d.conn, p, job, note=body.note if body else None)