448 lines
15 KiB
Python
448 lines
15 KiB
Python
"""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())
|