atc-agents/api/pipeline_ops.py

"""Pipeline / data-generation control for the Command Center.

Triggers per-database Airflow DAGs (light, configurable row counts), reports
run status, and returns live "in sync" counts via Trino so the UI can show the
whole data flow being pulsed and recognised downstream.
"""

from __future__ import annotations

import os
import time
from typing import Any

import httpx
from fastapi import APIRouter, Body, Query
from fastapi.responses import JSONResponse

AIRFLOW_URL = os.getenv("AIRFLOW_URL", "http://10.0.21.55:8080").rstrip("/")
AIRFLOW_USER = os.getenv("AIRFLOW_USER", "admin")
AIRFLOW_PASSWORD = os.getenv("AIRFLOW_PASSWORD", "")
TRINO_URL = os.getenv("TRINO_URL", "http://10.0.21.50:8089").rstrip("/")
TRINO_USER = os.getenv("TRINO_USER", "mo")

router = APIRouter(prefix="/api/pipeline", tags=["pipeline"])

# UI source key -> Airflow DAG id
SOURCE_DAG = {
    "postgres": "gen_postgres",
    "mysql": "gen_mysql",
    "mongodb": "gen_mongodb",
    "cassandra": "gen_cassandra",
    "neo4j": "gen_neo4j",
    "all": "generate_data_all_databases",
    "hadoop": "gen_hadoop_history",
}

# UI source key -> Trino fully-qualified table for live row counts
SOURCE_COUNT_SQL = {
    "postgres": "SELECT count(*) FROM postgres_sales.public.sales_orders",
    "mysql": "SELECT count(*) FROM mysql_hr.hr.employee_events",
    "mongodb": "SELECT count(*) FROM mongodb_supplychain.supplychain.events",
    "cassandra": "SELECT count(*) FROM cassandra_telemetry.telemetry.device_metrics",
}

_token_cache: dict[str, Any] = {"token": None, "exp": 0.0}


async def _airflow_token(client: httpx.AsyncClient) -> str:
    now = time.time()
    if _token_cache["token"] and _token_cache["exp"] > now + 30:
        return _token_cache["token"]
    r = await client.post(
        f"{AIRFLOW_URL}/auth/token",
        json={"username": AIRFLOW_USER, "password": AIRFLOW_PASSWORD},
        timeout=10,
    )
    r.raise_for_status()
    tok = r.json()["access_token"]
    _token_cache["token"] = tok
    _token_cache["exp"] = now + 20 * 60  # tokens last ~24h; refresh well before
    return tok


async def _trino_scalar(sql: str, timeout: float = 12.0) -> int | None:
    try:
        async with httpx.AsyncClient(timeout=timeout) as client:
            r = await client.post(
                f"{TRINO_URL}/v1/statement",
                content=sql.encode(),
                headers={"X-Trino-User": TRINO_USER},
            )
            data = r.json()
            rows: list[Any] = []
            nxt = data.get("nextUri")
            if data.get("data"):
                rows += data["data"]
            while nxt:
                rr = await client.get(nxt)
                d = rr.json()
                if d.get("data"):
                    rows += d["data"]
                if d.get("error"):
                    return None
                nxt = d.get("nextUri")
            if rows:
                return int(rows[0][0])
    except Exception:
        return None
    return None


@router.post("/generate/{source}")
async def generate(source: str, body: dict[str, Any] = Body(default={})) -> JSONResponse:
    dag_id = SOURCE_DAG.get(source)
    if not dag_id:
        return JSONResponse({"ok": False, "error": f"Unknown source '{source}'"}, status_code=400)
    rows = body.get("rows")
    conf: dict[str, Any] = {}
    if rows is not None:
        try:
            conf["rows"] = max(1, min(int(rows), 2_000_000))
        except (TypeError, ValueError):
            return JSONResponse({"ok": False, "error": "rows must be an integer"}, status_code=400)
    try:
        async with httpx.AsyncClient() as client:
            tok = await _airflow_token(client)
            r = await client.post(
                f"{AIRFLOW_URL}/api/v2/dags/{dag_id}/dagRuns",
                headers={"Authorization": f"Bearer {tok}"},
                json={"logical_date": None, "conf": conf},
                timeout=15,
            )
        if r.status_code >= 400:
            return JSONResponse({"ok": False, "error": f"Airflow {r.status_code}: {r.text[:300]}"}, status_code=200)
        j = r.json()
        return JSONResponse({
            "ok": True,
            "source": source,
            "dag_id": dag_id,
            "run_id": j.get("dag_run_id"),
            "state": j.get("state"),
            "rows": conf.get("rows"),
        })
    except Exception as exc:
        return JSONResponse({"ok": False, "error": str(exc)}, status_code=200)


@router.get("/runs/{source}")
async def runs(source: str, limit: int = Query(5)) -> JSONResponse:
    dag_id = SOURCE_DAG.get(source)
    if not dag_id:
        return JSONResponse({"ok": False, "error": f"Unknown source '{source}'"}, status_code=400)
    try:
        async with httpx.AsyncClient() as client:
            tok = await _airflow_token(client)
            headers = {"Authorization": f"Bearer {tok}"}
            r = await client.get(
                f"{AIRFLOW_URL}/api/v2/dags/{dag_id}/dagRuns",
                headers=headers,
                params={"order_by": "-run_after", "limit": limit},
                timeout=12,
            )
            runs_list = r.json().get("dag_runs", [])
            out = []
            for run in runs_list:
                out.append({
                    "run_id": run.get("dag_run_id"),
                    "state": run.get("state"),
                    "start": run.get("start_date"),
                    "end": run.get("end_date"),
                    "conf": run.get("conf"),
                })
        return JSONResponse({"ok": True, "source": source, "dag_id": dag_id, "runs": out})
    except Exception as exc:
        return JSONResponse({"ok": False, "error": str(exc)}, status_code=200)


@router.get("/sync")
async def sync(source: str | None = Query(None)) -> JSONResponse:
    """Live row counts per source (via Trino) for the 'in sync' display."""
    targets = [source] if source and source in SOURCE_COUNT_SQL else list(SOURCE_COUNT_SQL.keys())
    counts: dict[str, Any] = {}
    for src in targets:
        counts[src] = await _trino_scalar(SOURCE_COUNT_SQL[src])
    return JSONResponse({"ok": True, "counts": counts})