SupplyAgent/supply_infra/pipeline/step_cli.py

from __future__ import annotations

import argparse
import json
import logging
from pathlib import Path
from typing import Any

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 dispose_engine, get_session
from supply_infra.pipeline.contracts import StepContext
from supply_infra.pipeline.registry import execute_registered_step

logger = logging.getLogger(__name__)


def execute_step_run(step_run_id: str) -> dict[str, Any]:
    with get_session() as session:
        step = PipelineStepRunRepository(session).get(step_run_id)
        if step is None:
            raise ValueError(f"step_run_id not found: {step_run_id}")
        run = PipelineRunRepository(session).get(step.run_id)
        if run is None:
            raise ValueError(f"pipeline run not found: {step.run_id}")
        context = StepContext(
            run_id=run.run_id,
            step_run_id=step.step_run_id,
            step_key=step.step_key,
            biz_dt=run.biz_dt,
            date_snapshot=dict(run.date_snapshot_json or {}),
            config_snapshot=dict(run.config_snapshot_json or {}),
            input_snapshot=dict(step.input_snapshot_json or {}),
        )
    return execute_registered_step(context)


def main() -> None:
    parser = argparse.ArgumentParser(description="Execute one durable pipeline step")
    parser.add_argument("step_run_id")
    parser.add_argument("--result-file", required=True)
    args = parser.parse_args()
    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
    )
    result: dict[str, Any]
    exit_code = 0
    try:
        result = execute_step_run(args.step_run_id)
        if result.get("success") is False:
            exit_code = 1
    except Exception as exc:
        logger.exception("Pipeline step failed: step_run_id=%s", args.step_run_id)
        result = {"success": False, "error_code": "step_exception", "error": str(exc)}
        exit_code = 1
    finally:
        dispose_engine()
    Path(args.result_file).write_text(
        json.dumps(result, ensure_ascii=False, default=str),
        encoding="utf-8",
    )
    raise SystemExit(exit_code)


if __name__ == "__main__":
    main()