PacketStation/server/packetstation/worker/runs.py

"""One run of one job: advisory lock, job_runs row, request claiming, the job itself, the result.

`run_job` works with a started Scheduler (the worker) or an unstarted one (the CLI, the tests): it only needs
`settings`, `jobs`, `stack`, `host`, `build`, `connect()`, `set_feed()`, `running_add()`, `running_discard()`
and `after_run()`.
"""

from __future__ import annotations

import json
import logging
import traceback
from datetime import datetime
from typing import TYPE_CHECKING, Any

from psycopg.rows import dict_row
from psycopg.types.json import Jsonb

from .. import guard
from ..feed import config
from .jobs import JobContext, JobSpec
from .redaction import clip, redact

if TYPE_CHECKING:
    from .scheduler import Scheduler

log = logging.getLogger("packetstation.worker.runs")

LOCK_SPACE = 7310
FEED_JOB = "visual-feed"

__all__ = ["LOCK_SPACE", "redact", "run_job"]


def _jsonable(value: Any) -> Any:
    return json.loads(json.dumps(value if value is not None else {}, default=str))


def _row(conn: Any, run_id: int) -> dict[str, Any]:
    with conn.cursor(row_factory=dict_row) as cur:
        row = cur.execute("SELECT * FROM job_runs WHERE id = %s", (run_id,)).fetchone()
    return dict(row) if row else {"id": run_id}


def _insert(conn: Any, sch: Scheduler, name: str, trigger: str, *, status: str = "running", counts: Any = None) -> int:
    rid = conn.execute(
        "INSERT INTO job_runs (job, trigger, status, counts, build_version, host, finished_at)"
        " VALUES (%s, %s, %s, %s, %s, %s, CASE WHEN %s THEN now() END) RETURNING id",
        (name, trigger, status, Jsonb(_jsonable(counts)), sch.build, sch.host, status != "running"),
    ).fetchone()[0]
    return int(rid)


def _finish(conn: Any, run_id: int, *, status: str, counts: Any, errors: list[str], visual_target: str | None) -> None:
    conn.execute(
        "UPDATE job_runs SET status = %s, counts = %s, errors = %s, visual_target = coalesce(%s, visual_target),"
        " finished_at = now() WHERE id = %s",
        (status, Jsonb(_jsonable(counts)), Jsonb(clip(errors)), visual_target, run_id),
    )


def _target(result: Any) -> str:
    """ "<server> / <database>" from a guard result, without the (DEV) or (PRODUCTION) suffix."""
    text = str(getattr(result, "target", "") or "")
    return text.rsplit(" (", 1)[0] if text.endswith(")") else text


def run_job(sch: Scheduler, name: str, trigger: str) -> dict[str, Any]:
    """Run `name` once and return its job_runs row. Never raises for a job failure; raises when the database
    cannot be reached at all (the caller logs it)."""
    spec: JobSpec = sch.jobs[name]
    s = sch.settings
    conn = sch.connect()
    locked = False
    run_id: int | None = None
    status, counts, errors = "failed", {}, []  # type: tuple[str, Any, list[str]]
    visual_target: str | None = None
    pending_after = False
    sch.running_add(name)
    try:
        feed = config.load(conn)
        sch.set_feed(feed)
        locked = bool(conn.execute("SELECT pg_try_advisory_lock(%s, hashtext(%s))", (LOCK_SPACE, name)).fetchone()[0])
        if not locked:
            rid = _insert(conn, sch, name, trigger, status="skipped", counts={"reason": "already running"})
            log.info("job %s skipped: already running", name)
            return _row(conn, rid)

        with conn.transaction():
            run_id = _insert(conn, sch, name, trigger)
            claimed = conn.execute(
                "UPDATE job_requests SET claimed_run_id = %s, claimed_at = now()"
                " WHERE job = %s AND claimed_at IS NULL RETURNING requested_by",
                (run_id, name),
            ).fetchall()
            if claimed:
                trigger = "request"
                who = sorted({r[0] for r in claimed if r[0] is not None}, key=str)
                conn.execute("UPDATE job_runs SET trigger = 'request', requested_by = %s WHERE id = %s", (who, run_id))

        if trigger == "schedule" and spec.pausable and feed.paused:
            status, counts = "skipped", {"reason": "paused"}
        else:
            ctx = JobContext(
                conn=conn,
                settings=s,
                feed_settings=feed,
                run_id=run_id,
                trigger=trigger,
                started_at=datetime.now(sch.tz),
                log=logging.getLogger(f"packetstation.worker.jobs.{name}"),
                stack=sch.stack,
            )
            log.info("job %s started", name, extra={"ps": {"run_id": run_id, "trigger": trigger}})
            try:
                if name == FEED_JOB and s.feed_source == "visualkit":
                    # The stack guard again with fresh VisualKit settings: a test stack never reads production
                    ctx.visual_target = _target(guard.check_stack(s))
                counts = spec.run(ctx) or {}
                status = "degraded" if ctx.degraded else "ok"
            except Exception as e:
                status = "failed"
                ctx.errors.append(f"{type(e).__name__}: {e}")
                log.error("job %s failed: %s", name, redact(traceback.format_exc()))
            errors, visual_target = ctx.errors, ctx.visual_target
        try:
            _finish(conn, run_id, status=status, counts=counts, errors=errors, visual_target=visual_target)
        except Exception:
            # The run's connection is broken: write the result on a fresh one
            with sch.connect() as c2:
                _finish(c2, run_id, status=status, counts=counts, errors=errors, visual_target=visual_target)
        try:
            pending_after = bool(
                conn.execute(
                    "SELECT EXISTS (SELECT 1 FROM job_requests WHERE job = %s AND claimed_at IS NULL)", (name,)
                ).fetchone()[0]
            )
            row = _row(conn, run_id)
        except Exception:
            with sch.connect() as c2:
                row = _row(c2, run_id)
        log.info(
            "job %s finished: %s", name, status, extra={"ps": {"run_id": run_id, "status": status, "counts": counts}}
        )
        return row
    except Exception as e:
        status = "failed"
        if run_id is not None:
            try:
                with sch.connect() as c2:
                    _finish(
                        c2,
                        run_id,
                        status="failed",
                        counts=counts,
                        errors=[f"{type(e).__name__}: {e}"],
                        visual_target=visual_target,
                    )
            except Exception as e2:
                log.error("job %s: can't record the failure: %s", name, type(e2).__name__)
        raise
    finally:
        if locked:
            try:
                conn.execute("SELECT pg_advisory_unlock(%s, hashtext(%s))", (LOCK_SPACE, name))
            except Exception as e:
                log.warning("job %s: unlock failed (%s); closing the connection releases it", name, type(e).__name__)
        try:
            conn.close()
        finally:
            sch.running_discard(name)
            if run_id is not None:
                sch.after_run(name, status, pending_after)