PacketStation/server/packetstation/worker/jobs/__init__.py

"""The job registry and the shapes every job uses.

Jobs import `from packetstation.worker.jobs import JobContext, JobSpec, register` and expose `SPEC`. The feed and
drawing jobs are imported lazily by `load_jobs`, so the worker runs before they exist.
"""

from __future__ import annotations

import importlib
import logging
import threading
from collections.abc import Callable
from dataclasses import dataclass, field
from datetime import UTC, date, datetime, time, timedelta
from typing import Any
from zoneinfo import ZoneInfo

from ...feed.config import FeedSettings
from ...settings import Settings
from ..redaction import MAX_ERRORS, clip, redact

log = logging.getLogger(__name__)

Counts = dict[str, Any]
Cadence = Callable[[FeedSettings, datetime], int | None]
TRIGGERS = ("schedule", "request", "chain", "cli")


@dataclass
class Stack:
    """Per-process state shared by every run: the startup guard result and long-lived objects (the M-Files vault)."""

    guard: Any = None
    cache: dict[str, Any] = field(default_factory=dict)
    lock: threading.Lock = field(default_factory=threading.Lock)

    def cached(self, key: str, factory: Callable[[], Any]) -> Any:
        """The object under `key`, made once by `factory`. A factory that raises leaves nothing cached."""
        with self.lock:
            if key not in self.cache:
                self.cache[key] = factory()
            return self.cache[key]

    def drop(self, key: str) -> None:
        with self.lock:
            self.cache.pop(key, None)


@dataclass
class JobContext:
    """What a run gets. `conn` is the run's own autocommit connection as the worker role."""

    conn: Any
    settings: Settings
    feed_settings: FeedSettings
    run_id: int
    trigger: str
    started_at: datetime
    log: logging.Logger
    stack: Stack
    degraded: bool = False
    errors: list[str] = field(default_factory=list)
    visual_target: str | None = None

    def error(self, text: str, *, degrade: bool = True) -> None:
        """Record one error (redacted, clipped). It makes the run degraded unless degrade is False."""
        if len(self.errors) < MAX_ERRORS:
            self.errors.extend(clip([text]))
        if degrade:
            self.degraded = True

    def degrade(self, reason: str | None = None) -> None:
        self.degraded = True
        if reason:
            self.error(reason)

    @property
    def tz(self) -> ZoneInfo:
        return ZoneInfo(self.settings.site_tz)

    def now_site(self) -> datetime:
        return datetime.now(self.tz)

    def today_site(self) -> date:
        return self.now_site().date()


@dataclass(frozen=True)
class JobSpec:
    name: str
    run: Callable[[JobContext], Counts]
    cadence: Cadence  # seconds to the next run, or None for "no schedule"
    pausable: bool  # the feed pause skips scheduled fires (requests still run)
    chain: tuple[str, ...] = ()  # moved to now after a run that did not fail


JOBS: dict[str, JobSpec] = {}
_loaded = False
_load_lock = threading.Lock()
LAZY = ("visual_feed", "drawing_lookup")


def register(spec: JobSpec) -> JobSpec:
    JOBS[spec.name] = spec
    return spec


def load_jobs() -> dict[str, JobSpec]:
    """Register every job once and return the registry. A job module that cannot import is logged and left out."""
    global _loaded
    with _load_lock:
        if not _loaded:
            from . import housekeeping

            register(housekeeping.SPEC)
            for mod in LAZY:
                try:
                    m = importlib.import_module(f"{__name__}.{mod}")
                    register(m.SPEC)
                except ImportError as e:
                    log.warning("job module %s not loaded: %s", mod, redact(str(e)))
            _loaded = True
    return JOBS


def seconds_until(now_site: datetime, at: time) -> int:
    """Seconds from now_site to the next `at` (site wall clock), at least 1."""
    target = datetime.combine(now_site.date(), at, tzinfo=now_site.tzinfo)
    if target <= now_site:
        target = datetime.combine(now_site.date() + timedelta(days=1), at, tzinfo=now_site.tzinfo)
    # Aware datetimes sharing a tzinfo subtract as wall clock; through UTC a DST change counts right
    return max(1, int((target.astimezone(UTC) - now_site.astimezone(UTC)).total_seconds()))


__all__ = [
    "JOBS",
    "TRIGGERS",
    "Cadence",
    "Counts",
    "JobContext",
    "JobSpec",
    "Stack",
    "load_jobs",
    "register",
    "seconds_until",
]