tracer-cloud--opensre/integrations/tracer/tracer_pipelines.py

"""Pipeline and run-related API methods and models."""

from dataclasses import dataclass

from integrations.tracer.tracer_client_base import TracerClientBase


@dataclass(frozen=True)
class PipelineSummary:
    """Summary of a pipeline from the web app."""

    pipeline_name: str
    health_status: str | None
    last_run_start_time: str | None
    n_runs: int
    n_active_runs: int
    n_completed_runs: int


@dataclass(frozen=True)
class PipelineRunSummary:
    """Summary of a pipeline run from the web app."""

    pipeline_name: str
    run_id: str | None
    run_name: str | None
    trace_id: str | None
    status: str | None
    start_time: str | None
    end_time: str | None
    run_cost: float
    tool_count: int
    user_email: str | None
    instance_type: str | None
    region: str | None
    log_file_count: int


@dataclass(frozen=True)
class TracerRunResult:
    """Result from get_latest_run."""

    found: bool
    run_id: str | None = None
    pipeline_name: str | None = None
    run_name: str | None = None
    status: str | None = None
    start_time: str | None = None
    end_time: str | None = None
    run_time_seconds: float = 0
    run_cost: float = 0
    max_ram_gb: float = 0
    user_email: str | None = None
    team: str | None = None
    department: str | None = None
    instance_type: str | None = None
    environment: str | None = None
    region: str | None = None
    tool_count: int = 0


class TracerPipelinesMixin(TracerClientBase):
    """Mixin for Tracer pipeline and run-related API methods."""

    def get_pipelines(self, page: int = 1, size: int = 50) -> list[PipelineSummary]:
        """Fetch pipeline stats from /api/pipelines."""
        params = {"orgId": self.org_id, "page": page, "size": size}
        data = self._get("/api/pipelines", params)

        if not data.get("success") or not data.get("data"):
            return []

        pipelines = []
        for row in data["data"]:
            pipelines.append(
                PipelineSummary(
                    pipeline_name=row.get("pipeline_name", ""),
                    health_status=row.get("health_status"),
                    last_run_start_time=row.get("last_run_start_time"),
                    n_runs=int(row.get("n_runs", 0) or 0),
                    n_active_runs=int(row.get("n_active_runs", 0) or 0),
                    n_completed_runs=int(row.get("n_completed_runs", 0) or 0),
                )
            )
        return pipelines

    def get_pipeline_runs(
        self,
        pipeline_name: str,
        page: int = 1,
        size: int = 50,
    ) -> list[PipelineRunSummary]:
        """Fetch runs for a pipeline from /api/batch-runs."""
        params = {
            "orgId": self.org_id,
            "page": page,
            "size": size,
            "pipelineName": pipeline_name,
        }
        data = self._get("/api/batch-runs", params)

        if not data.get("success") or not data.get("data"):
            return []

        runs = []
        for row in data["data"]:
            runs.append(
                PipelineRunSummary(
                    pipeline_name=row.get("pipeline_name", pipeline_name),
                    run_id=row.get("run_id"),
                    run_name=row.get("run_name"),
                    trace_id=row.get("trace_id"),
                    status=row.get("status"),
                    start_time=row.get("start_time"),
                    end_time=row.get("end_time"),
                    run_cost=float(row.get("run_cost", 0) or 0),
                    tool_count=int(row.get("tool_count", 0) or 0),
                    user_email=row.get("user_email"),
                    instance_type=row.get("instance_type"),
                    region=row.get("region"),
                    log_file_count=int(row.get("log_file_count", 0) or 0),
                )
            )
        return runs

    def get_batch_details(self, trace_id: str) -> dict:
        """Fetch detailed batch run information from /api/batch-runs/[trace_id]."""
        params = {"orgId": self.org_id}
        data = self._get(f"/api/batch-runs/{trace_id}", params)
        return data

    def get_host_metrics(self, trace_id: str) -> dict:
        """Fetch host metrics (CPU, RAM, disk, GPU) from /api/runs/[trace_id]/host-metrics."""
        data = self._get(f"/api/runs/{trace_id}/host-metrics")
        return data

    def get_airflow_metrics(self, trace_id: str) -> dict:
        """Fetch Airflow metrics from /api/runs/[trace_id]/airflow."""
        params = {"orgId": self.org_id}
        data = self._get(f"/api/runs/{trace_id}/airflow", params)
        return data

    def get_latest_run(self, pipeline_name: str | None = None) -> TracerRunResult:
        """Get the latest (most recent) run for a pipeline from /api/batch-runs."""
        params: dict = {"page": 1, "size": 1, "orgId": self.org_id}
        if pipeline_name:
            params["pipelineName"] = pipeline_name
        data = self._get("/api/batch-runs", params)

        if not data.get("success") or not data.get("data"):
            return TracerRunResult(found=False)

        row = data["data"][0]
        tags = row.get("tags", {})
        max_ram_gb = float(row.get("max_ram", 0) or 0) / (1024**3)

        return TracerRunResult(
            found=True,
            run_id=row.get("run_id", ""),
            pipeline_name=row.get("pipeline_name", ""),
            run_name=row.get("run_name", ""),
            status=row.get("status", "Unknown"),
            start_time=row.get("start_time", ""),
            end_time=row.get("end_time"),
            run_time_seconds=float(row.get("run_time_seconds", 0) or 0),
            run_cost=float(row.get("run_cost", 0) or 0),
            max_ram_gb=max_ram_gb,
            user_email=tags.get("email", row.get("user_email", "")),
            team=tags.get("team", ""),
            department=tags.get("department", ""),
            instance_type=tags.get("instance_type", row.get("instance_type", "")),
            environment=row.get("environment", tags.get("environment", "")),
            region=row.get("region", ""),
            tool_count=int(row.get("tool_count", 0) or 0),
        )