SupplyAgent/api/services/pipeline.py

from __future__ import annotations

from typing import Any

from supply_infra.pipeline.dag import PIPELINE_STEPS
from supply_infra.pipeline.health import pipeline_health_snapshot
from supply_infra.pipeline.run_service import (
    cancel_pipeline_run,
    get_pipeline_run,
    list_pipeline_runs,
    resume_pipeline_run,
    retry_pipeline_step,
    submit_pipeline_run,
)


def create_run(
    *,
    biz_dt: str | None,
    source: str,
    reason: str | None = None,
    dry_run: bool = False,
) -> dict[str, Any]:
    return submit_pipeline_run(
        biz_dt=biz_dt,
        trigger_type="api",
        trigger_source=source,
        trigger_reason=reason,
        dry_run=dry_run,
    ).to_dict()


def list_runs(
    *,
    limit: int,
    status: str | None,
    biz_dt: str | None,
) -> list[dict[str, Any]]:
    return list_pipeline_runs(limit=limit, status=status, biz_dt=biz_dt)


def get_run(run_id: str) -> dict[str, Any] | None:
    return get_pipeline_run(run_id)


def resume_run(run_id: str) -> bool:
    return resume_pipeline_run(run_id)


def cancel_run(run_id: str) -> bool:
    return cancel_pipeline_run(run_id)


def retry_step(run_id: str, step_key: str) -> bool:
    return retry_pipeline_step(run_id, step_key)


def pipeline_health() -> dict[str, Any]:
    recent = list_pipeline_runs(limit=1)
    latest = recent[0] if recent else None
    return {
        "scheduler_embedded_in_api": False,
        "pipeline_key": "supply_pipeline",
        "steps": len(PIPELINE_STEPS),
        "latest_run": latest,
        **pipeline_health_snapshot(),
    }