feat: sync latest ARR implementation
This commit is contained in:
Wyndham ARR
447
monthly_reports/worker.py
Normal file
447
monthly_reports/worker.py
Normal file
@@ -0,0 +1,447 @@
|
||||
"""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())
|
||||
Reference in New Issue
Block a user