smart-org/apps/api/app/workers/base.py

"""Shared plumbing for the standalone workers.

The workers are deliberately thin. Every one of them is a *driver* that calls
the same application services the HTTP API calls — the wiring comes from
``app.application.container``, so a worker and the API can never disagree about
which adapter backs a port. No business rule lives in this package; if a worker
and the API ever disagreed, the API would be right.

That is also what makes the ``services/*`` split cheap later: moving a worker
into its own process changes the process boundary, not the code.
"""

from __future__ import annotations

import asyncio
import logging
from collections.abc import AsyncIterator, Awaitable, Callable
from contextlib import asynccontextmanager
from dataclasses import dataclass, field

from sqlalchemy.ext.asyncio import AsyncSession

from app.core.database import SessionFactory

logger = logging.getLogger("agy.worker")


@dataclass(slots=True)
class WorkerOutcome:
    """What one pass of a worker did.

    Returned rather than only logged, so a test can assert on it and an
    operator can see a number instead of grepping a log file.
    """

    name: str
    processed: int = 0
    skipped: int = 0
    failed: int = 0
    notes: list[str] = field(default_factory=list)

    #: Keep the note list bounded: a worker that processes 10k rows must not
    #: accumulate 10k strings in memory just to print a summary.
    MAX_NOTES = 50

    def note(self, message: str) -> None:
        if len(self.notes) < self.MAX_NOTES:
            self.notes.append(message)

    @property
    def ok(self) -> bool:
        return self.failed == 0

    @property
    def touched(self) -> int:
        return self.processed + self.skipped + self.failed

    def summary(self) -> str:
        return (
            f"{self.name}: processed={self.processed} "
            f"skipped={self.skipped} failed={self.failed}"
        )


@asynccontextmanager
async def worker_session() -> AsyncIterator[AsyncSession]:
    """One committed transaction per pass; rolled back whole on failure.

    A worker pass is all-or-nothing on purpose: a half-applied invoice rollup
    is worse than no rollup, because the next pass would skip the period it
    thinks it already billed.
    """
    async with SessionFactory() as session:
        try:
            yield session
            await session.commit()
        except Exception:
            await session.rollback()
            raise


async def run_periodically(
    pass_fn: Callable[[], Awaitable[WorkerOutcome]],
    *,
    interval_seconds: float,
    stop: asyncio.Event,
    run_once: bool = False,
    worker_name: str | None = None,
) -> list[WorkerOutcome]:
    """Drive ``pass_fn`` until ``stop`` is set — or exactly once.

    ``run_once=True`` is the cron/scheduled shape (what the CLI uses by
    default). The loop shape exists for long-running deployments.

    A failing pass is persisted before the loop continues. A worker that only
    logged to stderr would lose every failure the moment the container was
    replaced — which is precisely when someone wants to know what went wrong.
    """
    outcomes: list[WorkerOutcome] = []
    name: str = worker_name or str(getattr(pass_fn, "__name__", "worker"))
    while True:
        try:
            outcome = await pass_fn()
            outcomes.append(outcome)
            logger.info(outcome.summary())
            for note in outcome.notes:
                logger.info("%s | %s", outcome.name, note)
        except Exception as exc:  # noqa: BLE001 - a worker must survive a bad pass
            logger.exception("%s: pass failed", name)
            await _record_pass_failure(name, exc)
            outcomes.append(
                WorkerOutcome(name=name, failed=1, notes=[f"pass failed: {type(exc).__name__}"])
            )

        if run_once:
            return outcomes

        try:
            # Returns early (cleanly) when ``stop`` is set during the wait.
            await asyncio.wait_for(stop.wait(), timeout=interval_seconds)
            return outcomes
        except TimeoutError:
            continue


async def _record_pass_failure(worker_name: str, exc: BaseException) -> None:
    """Persist a failed pass. Imported lazily to keep this module import-light."""
    from app.domain.audit.models import ErrorSeverity, ErrorSource
    from app.infrastructure.error_log import record_error

    await record_error(
        kind="worker_pass_failed",
        message=str(exc) or type(exc).__name__,
        exc=exc,
        source=ErrorSource.WORKER,
        severity=ErrorSeverity.ERROR,
        metadata={"worker": worker_name},
    )


def install_signal_handlers(stop: asyncio.Event) -> None:
    """Ask the event loop to set ``stop`` on SIGINT/SIGTERM.

    Best-effort: on a platform without these signals (or outside the main
    thread) the worker simply runs until interrupted some other way.
    """
    loop = asyncio.get_running_loop()
    for signal_name in ("SIGINT", "SIGTERM"):
        try:
            import signal as _signal

            sig = getattr(_signal, signal_name)
        except (AttributeError, ValueError):
            continue
        try:
            loop.add_signal_handler(sig, stop.set)
        except (NotImplementedError, RuntimeError, ValueError):
            continue


__all__ = [
    "WorkerOutcome",
    "install_signal_handlers",
    "run_periodically",
    "worker_session",
]