SupplyAgent/supply_infra/db/models/pipeline_run.py

from __future__ import annotations

from datetime import datetime
from typing import Any

from sqlalchemy import Boolean, DateTime, Index, JSON, String, Text, UniqueConstraint, func
from sqlalchemy.orm import Mapped, mapped_column

from supply_infra.db.base import Base


class PipelineRun(Base):
    """One durable execution request for the complete supply pipeline."""

    __tablename__ = "pipeline_run"
    __table_args__ = (
        UniqueConstraint("dedupe_key", name="uk_pipeline_run_dedupe_key"),
        Index("idx_pipeline_run_biz", "pipeline_key", "biz_dt", "created_at"),
        Index("idx_pipeline_run_status", "status", "lease_until"),
        Index("idx_pipeline_run_deadline", "status", "deadline_at"),
    )

    run_id: Mapped[str] = mapped_column(String(36), primary_key=True)
    dedupe_key: Mapped[str] = mapped_column(String(191), nullable=False)
    pipeline_key: Mapped[str] = mapped_column(
        String(64),
        nullable=False,
        default="supply_pipeline",
    )
    biz_dt: Mapped[str] = mapped_column(String(8), nullable=False)
    trigger_type: Mapped[str] = mapped_column(String(24), nullable=False)
    trigger_source: Mapped[str | None] = mapped_column(String(64), nullable=True)
    trigger_reason: Mapped[str | None] = mapped_column(Text, nullable=True)
    run_mode: Mapped[str] = mapped_column(String(24), nullable=False, default="full")
    dry_run: Mapped[bool] = mapped_column(Boolean, nullable=False, default=False)
    status: Mapped[str] = mapped_column(String(32), nullable=False, default="queued")
    current_step: Mapped[str | None] = mapped_column(String(64), nullable=True)
    deadline_at: Mapped[datetime | None] = mapped_column(DateTime, nullable=True)
    started_at: Mapped[datetime | None] = mapped_column(DateTime, nullable=True)
    finished_at: Mapped[datetime | None] = mapped_column(DateTime, nullable=True)
    heartbeat_at: Mapped[datetime | None] = mapped_column(DateTime, nullable=True)
    lease_owner: Mapped[str | None] = mapped_column(String(191), nullable=True)
    lease_until: Mapped[datetime | None] = mapped_column(DateTime, nullable=True)
    code_version: Mapped[str | None] = mapped_column(String(128), nullable=True)
    config_snapshot_json: Mapped[dict[str, Any]] = mapped_column(
        JSON,
        nullable=False,
        default=dict,
    )
    date_snapshot_json: Mapped[dict[str, Any]] = mapped_column(
        JSON,
        nullable=False,
        default=dict,
    )
    summary_json: Mapped[dict[str, Any] | None] = mapped_column(JSON, nullable=True)
    error_code: Mapped[str | None] = mapped_column(String(64), nullable=True)
    error_message: Mapped[str | None] = mapped_column(Text, nullable=True)
    created_at: Mapped[datetime] = mapped_column(
        DateTime,
        nullable=False,
        server_default=func.now(),
    )
    updated_at: Mapped[datetime] = mapped_column(
        DateTime,
        nullable=False,
        server_default=func.now(),
        onupdate=func.now(),
    )