SupplyAgent/supply_infra/pipeline/run_service.py

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,
    dry_run: bool = False,
    settings: InfraSettings | None = None,
) -> RunSubmission:
    active = settings or get_infra_settings()
    resolved_biz_dt = resolve_biz_dt(biz_dt, settings=active)
    run_mode = "dry_run" if dry_run else "full"
    dedupe_key = f"{PIPELINE_KEY}:{resolved_biz_dt}:{run_mode}"
    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=bool(existing.dry_run),
                )

            # 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": run_mode,
                    "dry_run": dry_run,
                    "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=dry_run,
            )
    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=bool(existing.dry_run),
            )


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,
                    "attempt_count": item.attempt_count,
                    "last_attempt_at": _iso(item.last_attempt_at),
                    "next_retry_at": _iso(item.next_retry_at),
                    "external_id": item.external_id,
                    "external_name": item.external_name,
                    "error": item.record_error,
                    "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()