SupplyAgent/supply_infra/pipeline/health.py

from __future__ import annotations

from datetime import datetime, timedelta
from typing import Any

from supply_infra.config import InfraSettings, get_infra_settings
from supply_infra.db.repositories.pipeline_lock_repo import PipelineLockRepository
from supply_infra.db.repositories.pipeline_run_repo import PipelineRunRepository
from supply_infra.db.session import get_session
from supply_infra.pipeline.dates import china_now

SCHEDULER_HEARTBEAT_KEY = "pipeline:scheduler"
_ACTIVE_STATUSES = {
    "queued",
    "running",
    "cancelling",
    "blocked_previous_run",
}


def touch_scheduler_heartbeat(
    *,
    owner: str,
    settings: InfraSettings | None = None,
) -> None:
    active = settings or get_infra_settings()
    with get_session() as session:
        PipelineLockRepository(session).touch(
            lock_key=SCHEDULER_HEARTBEAT_KEY,
            owner=owner,
            lease_seconds=max(
                active.pipeline_lease_seconds,
                active.pipeline_heartbeat_seconds * 3,
            ),
            now=china_now(),
        )


def run_alert_level(
    *,
    status: str,
    started_at: datetime | None,
    created_at: datetime,
    deadline_at: datetime | None,
    now: datetime,
    settings: InfraSettings,
) -> str | None:
    if status not in _ACTIVE_STATUSES:
        return None
    if deadline_at is not None:
        remaining = deadline_at - now
        if remaining <= timedelta(minutes=settings.pipeline_final_warn_before_next_minutes):
            return "final_warning"
        if remaining <= timedelta(minutes=settings.pipeline_critical_before_next_minutes):
            return "critical"
    baseline = started_at or created_at
    if now - baseline >= timedelta(hours=settings.pipeline_warn_after_hours):
        return "warning"
    return None


def pipeline_health_snapshot(
    settings: InfraSettings | None = None,
) -> dict[str, Any]:
    active = settings or get_infra_settings()
    now = china_now()
    with get_session() as session:
        scheduler_lock = PipelineLockRepository(session).get(SCHEDULER_HEARTBEAT_KEY)
        latest = PipelineRunRepository(session).list_recent(limit=1)
        latest_run = latest[0] if latest else None
        scheduler_heartbeat_at = (
            scheduler_lock.heartbeat_at if scheduler_lock is not None else None
        )
        scheduler_lease_until = (
            scheduler_lock.lease_until if scheduler_lock is not None else None
        )
        latest_snapshot = (
            {
                "run_id": latest_run.run_id,
                "status": latest_run.status,
                "started_at": latest_run.started_at,
                "created_at": latest_run.created_at,
                "deadline_at": latest_run.deadline_at,
            }
            if latest_run is not None
            else None
        )

    scheduler_heartbeat = (
        scheduler_heartbeat_at.isoformat()
        if scheduler_heartbeat_at is not None
        else None
    )
    scheduler_running = bool(
        active.scheduler_enabled
        and scheduler_lease_until is not None
        and scheduler_lease_until > now
    )
    alert_level = None
    if latest_snapshot is not None:
        alert_level = run_alert_level(
            status=latest_snapshot["status"],
            started_at=latest_snapshot["started_at"],
            created_at=latest_snapshot["created_at"],
            deadline_at=latest_snapshot["deadline_at"],
            now=now,
            settings=active,
        )
    return {
        "scheduler_enabled": active.scheduler_enabled,
        "scheduler_running": scheduler_running,
        "scheduler_heartbeat_at": scheduler_heartbeat,
        "latest_run_id": latest_snapshot["run_id"] if latest_snapshot else None,
        "latest_run_status": latest_snapshot["status"] if latest_snapshot else None,
        "latest_run_alert_level": alert_level,
        "checked_at": now.isoformat(),
    }