"""Dedicated reliable outbox worker for automatic monthly publication.""" from __future__ import annotations import argparse import json import signal import sys import time from dataclasses import dataclass from pathlib import Path from typing import Any, Callable, Mapping, Optional, Protocol, Sequence from arr_database import controlled_connect from monthly_reports.contracts import ErrorCode from monthly_reports.publishing import ArtifactToolBuilder, AtomicReportPublisher from monthly_reports.repository import ( DatabaseConfig, DerivedMonthlyRequest, PostgresReportRepository, RepositoryError, ) from monthly_reports.service import MonthlyReportService, RunRequest, RunResult PROJECT_ROOT = Path(__file__).resolve().parents[1] EVENT_TYPE = "arr.daily_version_committed" DEFAULT_LEASE_SECONDS = 300 DEFAULT_MAX_ATTEMPTS = 8 CLAIM_EVENT_SQL = """ WITH candidate AS ( SELECT event.id FROM ingestion.outbox_events AS event WHERE event.event_type = %s AND event.aggregate_type = 'processing_run' AND event.available_at <= now() AND event.publish_status IN ('pending', 'publishing') ORDER BY event.created_at, event.id FOR UPDATE SKIP LOCKED LIMIT 1 ) UPDATE ingestion.outbox_events AS event SET publish_status = 'publishing', publish_attempts = event.publish_attempts + 1, available_at = now() + (%s * interval '1 second'), last_error_code = NULL FROM candidate WHERE event.id = candidate.id RETURNING event.id, event.event_key, event.payload, event.publish_attempts """.strip() PUBLISH_EVENT_SQL = """ UPDATE ingestion.outbox_events AS event SET publish_status = 'published', published_at = now(), available_at = now(), last_error_code = NULL WHERE event.id = %s AND event.publish_status = 'publishing' AND EXISTS ( SELECT 1 FROM reporting.monthly_runs AS run JOIN ingestion.artifacts AS workbook ON workbook.id = run.workbook_artifact_id JOIN ingestion.artifacts AS result ON result.id = run.result_artifact_id WHERE run.id = %s AND run.report_status IN ('active', 'superseded') AND workbook.artifact_kind = 'monthly_xlsx' AND result.artifact_kind = 'result_json' ) """.strip() FAIL_EVENT_SQL = """ UPDATE ingestion.outbox_events SET publish_status = %s, available_at = CASE WHEN %s = 'pending' THEN now() + (%s * interval '1 second') ELSE now() END, published_at = NULL, last_error_code = %s WHERE id = %s AND publish_status = 'publishing' """.strip() class WorkerError(RuntimeError): def __init__(self, code: str, *, retryable: bool): super().__init__(code) self.code = code self.retryable = retryable @dataclass(frozen=True) class OutboxEvent: event_id: int event_key: str payload: Mapping[str, Any] publish_attempts: int @dataclass(frozen=True) class WorkerOutcome: status: str event_id: Optional[int] = None report_id: Optional[int] = None error_code: Optional[str] = None def to_dict(self) -> dict[str, Any]: payload: dict[str, Any] = {"status": self.status} if self.event_id is not None: payload["event_id"] = self.event_id if self.report_id is not None: payload["report_id"] = self.report_id if self.error_code is not None: payload["error_code"] = self.error_code return payload class OutboxRepository(Protocol): def claim_next(self) -> Optional[OutboxEvent]: ... def mark_published(self, event_id: int, report_id: int) -> None: ... def mark_failed(self, event: OutboxEvent, code: str, *, retryable: bool) -> str: ... class MonthlyRequestRepository(Protocol): def derive_monthly_request( self, daily_version_id: int, ) -> Optional[DerivedMonthlyRequest]: ... def _default_connect(dsn: str) -> Any: try: import psycopg # type: ignore[import-not-found] except ImportError: raise WorkerError("MONTHLY_WORKER_DATABASE_DRIVER_UNAVAILABLE", retryable=False) from None try: return psycopg.connect(dsn, autocommit=False) except Exception: raise WorkerError("MONTHLY_WORKER_DATABASE_UNAVAILABLE", retryable=True) from None class PostgresOutboxRepository: def __init__( self, config: DatabaseConfig, *, connect: Optional[Callable[[str], Any]] = None, lease_seconds: int = DEFAULT_LEASE_SECONDS, max_attempts: int = DEFAULT_MAX_ATTEMPTS, ) -> None: if lease_seconds < 30 or lease_seconds > 3600: raise ValueError("worker lease must be between 30 and 3600 seconds") if max_attempts < 1 or max_attempts > 100: raise ValueError("worker max attempts is invalid") self._config = config self._connect = connect or _default_connect self._lease_seconds = lease_seconds self._max_attempts = max_attempts def _open(self) -> Any: try: return self._connect(self._config.dsn) except WorkerError: raise except Exception: raise WorkerError("MONTHLY_WORKER_DATABASE_UNAVAILABLE", retryable=True) from None @staticmethod def _begin(cursor: Any) -> None: cursor.execute("SET TRANSACTION ISOLATION LEVEL READ COMMITTED") cursor.execute("SET LOCAL lock_timeout = '10s'") cursor.execute("SET LOCAL statement_timeout = '30s'") cursor.execute("SELECT current_database()") row = cursor.fetchone() if not row or row[0] != "booking_test": raise WorkerError("MONTHLY_WORKER_DATABASE_TARGET_INVALID", retryable=False) def claim_next(self) -> Optional[OutboxEvent]: connection = self._open() try: with connection.transaction(): with connection.cursor() as cursor: self._begin(cursor) cursor.execute(CLAIM_EVENT_SQL, (EVENT_TYPE, self._lease_seconds)) row = cursor.fetchone() if not row: return None payload = row[2] if isinstance(payload, str): payload = json.loads(payload) if not isinstance(payload, Mapping): payload = {} return OutboxEvent( event_id=int(row[0]), event_key=str(row[1]), payload=payload, publish_attempts=int(row[3]), ) except WorkerError: raise except Exception: raise WorkerError("MONTHLY_WORKER_CLAIM_FAILED", retryable=True) from None finally: connection.close() def mark_published(self, event_id: int, report_id: int) -> None: connection = self._open() try: with connection.transaction(): with connection.cursor() as cursor: self._begin(cursor) cursor.execute(PUBLISH_EVENT_SQL, (event_id, report_id)) if cursor.rowcount != 1: raise WorkerError("MONTHLY_WORKER_PUBLISH_ACK_FAILED", retryable=True) except WorkerError: raise except Exception: raise WorkerError("MONTHLY_WORKER_PUBLISH_ACK_FAILED", retryable=True) from None finally: connection.close() def mark_failed(self, event: OutboxEvent, code: str, *, retryable: bool) -> str: terminal = not retryable or event.publish_attempts >= self._max_attempts status = "dead" if terminal else "pending" delay_seconds = min(300, 2 ** min(event.publish_attempts, 8)) safe_code = code if code and len(code) <= 128 else "MONTHLY_WORKER_FAILED" connection = self._open() try: with connection.transaction(): with connection.cursor() as cursor: self._begin(cursor) cursor.execute( FAIL_EVENT_SQL, ( status, status, delay_seconds, safe_code, event.event_id, ), ) if cursor.rowcount != 1: raise WorkerError("MONTHLY_WORKER_FAILURE_ACK_FAILED", retryable=True) return status except WorkerError: raise except Exception: raise WorkerError("MONTHLY_WORKER_FAILURE_ACK_FAILED", retryable=True) from None finally: connection.close() class MonthlyOutboxWorker: def __init__( self, outbox: OutboxRepository, requests: MonthlyRequestRepository, service: MonthlyReportService, ) -> None: self._outbox = outbox self._requests = requests self._service = service @staticmethod def _daily_version_id(event: OutboxEvent) -> int: value = event.payload.get("daily_version_id") if isinstance(value, bool): raise WorkerError("MONTHLY_WORKER_EVENT_INVALID", retryable=False) try: daily_version_id = int(value) except (TypeError, ValueError): raise WorkerError("MONTHLY_WORKER_EVENT_INVALID", retryable=False) from None if daily_version_id < 1: raise WorkerError("MONTHLY_WORKER_EVENT_INVALID", retryable=False) return daily_version_id @staticmethod def _result_failure(result: RunResult) -> WorkerError: code = result.error_code or "MONTHLY_WORKER_REPORT_FAILED" retryable = code not in { ErrorCode.REQUEST_INVALID, ErrorCode.SOURCE_INVALID, } return WorkerError(code, retryable=retryable) def process_next(self) -> WorkerOutcome: event = self._outbox.claim_next() if event is None: return WorkerOutcome(status="idle") try: daily_version_id = self._daily_version_id(event) derived = self._requests.derive_monthly_request(daily_version_id) if derived is None: raise WorkerError("MONTHLY_WORKER_ARRIVAL_SCOPE_EMPTY", retryable=False) result = self._service.run( RunRequest( derived.report_year, derived.report_month, derived.as_of_date, ) ) if result.status != "success": raise self._result_failure(result) if ( result.report_version_id is None or result.artifact is None or not result.artifact.get("sha256") ): raise WorkerError("MONTHLY_WORKER_REPORT_RECEIPT_INVALID", retryable=True) self._outbox.mark_published(event.event_id, result.report_version_id) return WorkerOutcome( status="published", event_id=event.event_id, report_id=result.report_version_id, ) except WorkerError as error: status = self._outbox.mark_failed( event, error.code, retryable=error.retryable, ) return WorkerOutcome( status=status, event_id=event.event_id, error_code=error.code, ) except RepositoryError as error: status = self._outbox.mark_failed(event, error.code, retryable=True) return WorkerOutcome( status=status, event_id=event.event_id, error_code=error.code, ) except Exception: code = "MONTHLY_WORKER_INTERNAL_ERROR" status = self._outbox.mark_failed(event, code, retryable=True) return WorkerOutcome( status=status, event_id=event.event_id, error_code=code, ) def _parser() -> argparse.ArgumentParser: parser = argparse.ArgumentParser(prog="arr-monthly-worker") parser.add_argument("--db-config", type=Path) parser.add_argument("--driver-path", type=Path) parser.add_argument("--node-binary", type=Path) parser.add_argument("--artifact-tool-module", type=Path) parser.add_argument("--output-root", type=Path) parser.add_argument("--poll-seconds", type=float, default=2.0) parser.add_argument("--lease-seconds", type=int, default=DEFAULT_LEASE_SECONDS) parser.add_argument("--max-attempts", type=int, default=DEFAULT_MAX_ATTEMPTS) parser.add_argument("--once", action="store_true") return parser def _runtime(args: argparse.Namespace) -> MonthlyOutboxWorker: connect: Optional[Callable[[str], Any]] = None if args.db_config: connect = controlled_connect(args.db_config, args.driver_path) config = DatabaseConfig("controlled") else: config = DatabaseConfig.from_environment() repository = PostgresReportRepository(config, connect=connect) outbox = PostgresOutboxRepository( config, connect=connect, lease_seconds=args.lease_seconds, max_attempts=args.max_attempts, ) output_root = (args.output_root or PROJECT_ROOT / "outputs" / "monthly_reports").resolve() try: output_root.relative_to(PROJECT_ROOT) except ValueError: raise ValueError("worker output root must be inside the project") from None builder = ArtifactToolBuilder( PROJECT_ROOT / "monthly_reports" / "xlsx" / "build_workbook.mjs", node_binary=str(args.node_binary) if args.node_binary else None, artifact_tool_module=args.artifact_tool_module, ) service = MonthlyReportService( repository, builder, AtomicReportPublisher(PROJECT_ROOT, output_root), output_root / ".staging", ) return MonthlyOutboxWorker(outbox, repository, service) def _emit(outcome: WorkerOutcome) -> None: sys.stdout.write(json.dumps(outcome.to_dict(), sort_keys=True) + "\n") sys.stdout.flush() def main(argv: Optional[Sequence[str]] = None) -> int: args = _parser().parse_args(argv) if args.poll_seconds < 0.1 or args.poll_seconds > 60: raise SystemExit("poll interval must be between 0.1 and 60 seconds") worker = _runtime(args) if args.once: try: outcome = worker.process_next() except WorkerError as error: outcome = WorkerOutcome(status="worker_error", error_code=error.code) _emit(outcome) return 0 if outcome.status in {"idle", "published"} else 2 stopping = False def stop(_signum: int, _frame: object) -> None: nonlocal stopping stopping = True signal.signal(signal.SIGTERM, stop) signal.signal(signal.SIGINT, stop) while not stopping: try: outcome = worker.process_next() except WorkerError as error: outcome = WorkerOutcome(status="worker_error", error_code=error.code) if outcome.status != "idle": _emit(outcome) if outcome.status in {"idle", "worker_error"} and not stopping: time.sleep(args.poll_seconds) return 0 if __name__ == "__main__": raise SystemExit(main())