from __future__ import annotations import os import uuid from dataclasses import dataclass from datetime import datetime from typing import Any from sqlalchemy.exc import IntegrityError from supply_infra.config import InfraSettings, get_infra_settings from supply_infra.db.models.pipeline_run import PipelineRun from supply_infra.db.models.pipeline_step_run import PipelineStepRun from supply_infra.db.repositories.pipeline_lock_repo import PipelineLockRepository from supply_infra.db.repositories.pipeline_outbox_repo import PipelineOutboxRepository from supply_infra.db.repositories.pipeline_run_repo import PipelineRunRepository from supply_infra.db.repositories.pipeline_step_run_repo import PipelineStepRunRepository from supply_infra.db.session import get_session from supply_infra.pipeline.dag import PIPELINE_STEPS from supply_infra.pipeline.dates import ( build_date_snapshot, CHINA_TIMEZONE, china_now, deadline_for_trigger, next_schedule_china, resolve_biz_dt, ) from supply_infra.pipeline.enums import RunStatus, StepStatus PIPELINE_KEY = "supply_pipeline" @dataclass(frozen=True) class RunSubmission: run_id: str created: bool status: str biz_dt: str dry_run: bool = False def to_dict(self) -> dict[str, Any]: return { "accepted": True, "run_id": self.run_id, "created": self.created, "status": self.status, "biz_dt": self.biz_dt, "dry_run": self.dry_run, } def _config_snapshot(settings: InfraSettings) -> dict[str, Any]: return { "scheduler_timezone": settings.scheduler_timezone, "scheduler_cron_hour": settings.scheduler_cron_hour, "scheduler_cron_minute": settings.scheduler_cron_minute, "lease_seconds": settings.pipeline_lease_seconds, "heartbeat_seconds": settings.pipeline_heartbeat_seconds, "worker_processes": settings.pipeline_worker_processes, "max_active_steps": settings.pipeline_max_active_steps, "mysql_connection_budget": settings.mysql_connection_budget, } def _code_version() -> str | None: return os.getenv("BUILD_TIMESTAMP") or os.getenv("GIT_COMMIT") def submit_pipeline_run( *, biz_dt: str | None = None, trigger_type: str = "api", trigger_source: str | None = None, trigger_reason: str | None = None, settings: InfraSettings | None = None, ) -> RunSubmission: active = settings or get_infra_settings() resolved_biz_dt = resolve_biz_dt(biz_dt, settings=active) dedupe_key = f"{PIPELINE_KEY}:{resolved_biz_dt}:full" run_id = str(uuid.uuid4()) date_snapshot = build_date_snapshot(resolved_biz_dt) try: with get_session() as session: run_repo = PipelineRunRepository(session) PipelineLockRepository(session).lock_control_plane(now=china_now()) existing = run_repo.get_by_dedupe_key(dedupe_key) if existing is not None: return RunSubmission( run_id=existing.run_id, created=False, status=existing.status, biz_dt=existing.biz_dt, dry_run=False, ) # All entry points share the same no-overlap rule. Linking each new # run to the latest nonterminal run also creates a safe FIFO chain # when several future batches are submitted in advance. previous = run_repo.get_latest_nonterminal() initial_status = ( RunStatus.BLOCKED_PREVIOUS_RUN.value if previous is not None else RunStatus.QUEUED.value ) run = run_repo.create( { "run_id": run_id, "dedupe_key": dedupe_key, "pipeline_key": PIPELINE_KEY, "biz_dt": resolved_biz_dt, "trigger_type": trigger_type, "trigger_source": trigger_source or trigger_type, "trigger_reason": trigger_reason, "run_mode": "full", "dry_run": False, "status": initial_status, "deadline_at": deadline_for_trigger( trigger_type, settings=active, ), "code_version": _code_version(), "config_snapshot_json": _config_snapshot(active), "date_snapshot_json": date_snapshot, "summary_json": ( {"blocked_by_run_id": previous.run_id} if previous is not None else None ), } ) rows: list[dict[str, Any]] = [] for index, step in enumerate(PIPELINE_STEPS): if previous is not None: status = StepStatus.BLOCKED.value else: status = ( StepStatus.READY.value if index == 0 else StepStatus.PENDING.value ) rows.append( { "step_run_id": str(uuid.uuid4()), "run_id": run.run_id, "step_key": step.key, "step_order": step.order, "attempt": 1, "status": status, "critical": step.critical, "dependency_snapshot_json": list(step.dependencies), "input_snapshot_json": {}, "timeout_seconds": step.timeout_seconds, "max_attempts": step.max_attempts, "retryable": step.retryable, "error_code": ( "previous_run_active" if previous is not None else None ), "error_message": ( f"Blocked by previous run {previous.run_id}" if previous is not None else None ), } ) PipelineStepRunRepository(session).create_steps(rows) return RunSubmission( run_id=run.run_id, created=True, status=run.status, biz_dt=run.biz_dt, dry_run=False, ) except IntegrityError: with get_session() as session: existing = PipelineRunRepository(session).get_by_dedupe_key(dedupe_key) if existing is None: raise return RunSubmission( run_id=existing.run_id, created=False, status=existing.status, biz_dt=existing.biz_dt, dry_run=False, ) def get_pipeline_run(run_id: str) -> dict[str, Any] | None: with get_session() as session: run = PipelineRunRepository(session).get(run_id) if run is None: return None steps = PipelineStepRunRepository(session).list_for_run(run_id) outbox = PipelineOutboxRepository(session).list_for_run(run_id) return { **serialize_run(run), "steps": [serialize_step(item) for item in steps], "effects": [ { "outbox_id": item.outbox_id, "step_run_id": item.step_run_id, "effect_type": item.effect_type, "payload_hash": item.payload_hash, "payload_uri": item.payload_uri, "status": item.status, "dry_run": item.dry_run, "created_at": _iso(item.created_at), } for item in outbox ], } def list_pipeline_runs( *, limit: int = 50, status: str | None = None, biz_dt: str | None = None, ) -> list[dict[str, Any]]: with get_session() as session: rows = PipelineRunRepository(session).list_recent( limit=min(max(limit, 1), 200), status=status, biz_dt=biz_dt, ) return [serialize_run(item) for item in rows] def cancel_pipeline_run(run_id: str) -> bool: now = china_now() with get_session() as session: run_repo = PipelineRunRepository(session) step_repo = PipelineStepRunRepository(session) run = run_repo.get(run_id, for_update=True) if run is None: return False if run.status in { RunStatus.SUCCEEDED.value, RunStatus.FAILED.value, RunStatus.CANCELLED.value, RunStatus.DEADLINE_EXCEEDED.value, }: return True has_running_step = False for step in step_repo.list_for_run(run_id): if step.status == StepStatus.RUNNING.value: has_running_step = True continue if step.status in { StepStatus.PENDING.value, StepStatus.READY.value, StepStatus.RETRY_WAIT.value, StepStatus.BLOCKED.value, }: step.status = StepStatus.CANCELLED.value step.finished_at = now if has_running_step: run.status = RunStatus.CANCELLING.value run.error_code = "cancellation_requested" run.error_message = "Waiting for the running step to stop" run.heartbeat_at = now run.lease_owner = None run.lease_until = None else: run_repo.finish(run, status=RunStatus.CANCELLED.value, now=now) return True def resume_pipeline_run(run_id: str, *, step_key: str | None = None) -> bool: with get_session() as session: run_repo = PipelineRunRepository(session) step_repo = PipelineStepRunRepository(session) run = run_repo.get(run_id, for_update=True) if run is None: return False if run.status not in { RunStatus.FAILED.value, RunStatus.CANCELLED.value, RunStatus.DEADLINE_EXCEEDED.value, }: return False latest = step_repo.latest_attempts(run_id) resumable_statuses = ( {StepStatus.CANCELLED.value} if run.status == RunStatus.CANCELLED.value else { StepStatus.FAILED.value, StepStatus.TIMED_OUT.value, StepStatus.BLOCKED.value, } ) resume_from = next( ( item for item in sorted(latest.values(), key=lambda row: row.step_order) if item.status in resumable_statuses ), None, ) if resume_from is None: return False if step_key is not None and resume_from.step_key != step_key: return False retry = step_repo.create_retry_attempt( resume_from, status=StepStatus.READY.value, ) for item in latest.values(): if item.step_order > resume_from.step_order and item.status in { StepStatus.BLOCKED.value, StepStatus.CANCELLED.value, }: item.status = StepStatus.PENDING.value item.error_code = None item.error_message = None run.status = RunStatus.QUEUED.value run.current_step = retry.step_key if run.deadline_at is not None and run.deadline_at <= china_now(): run.deadline_at = next_schedule_china() run.finished_at = None run.summary_json = None run.error_code = None run.error_message = None return True def retry_pipeline_step(run_id: str, step_key: str) -> bool: return resume_pipeline_run(run_id, step_key=step_key) def serialize_run(run: PipelineRun) -> dict[str, Any]: return { "run_id": run.run_id, "pipeline_key": run.pipeline_key, "biz_dt": run.biz_dt, "trigger_type": run.trigger_type, "trigger_source": run.trigger_source, "run_mode": run.run_mode, "dry_run": run.dry_run, "status": run.status, "current_step": run.current_step, "deadline_at": _iso(run.deadline_at), "started_at": _iso(run.started_at), "finished_at": _iso(run.finished_at), "heartbeat_at": _iso(run.heartbeat_at), "summary": run.summary_json, "error_code": run.error_code, "error_message": run.error_message, "created_at": _iso(run.created_at), "updated_at": _iso(run.updated_at), } def serialize_step(step: PipelineStepRun) -> dict[str, Any]: return { "step_run_id": step.step_run_id, "run_id": step.run_id, "step_key": step.step_key, "step_order": step.step_order, "attempt": step.attempt, "status": step.status, "critical": step.critical, "timeout_seconds": step.timeout_seconds, "started_at": _iso(step.started_at), "finished_at": _iso(step.finished_at), "heartbeat_at": _iso(step.heartbeat_at), "result": step.result_summary_json, "error_code": step.error_code, "error_message": step.error_message, "log_uri": step.log_uri, } def _iso(value: datetime | None) -> str | None: if value is None: return None aware = ( value.replace(tzinfo=CHINA_TIMEZONE) if value.tzinfo is None else value.astimezone(CHINA_TIMEZONE) ) return aware.isoformat()