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 asyncio
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 -> the agent responsible for that part of the platform
SOURCE_AGENT = {
    "postgres": "data-custodian",
    "mysql": "data-custodian",
    "mongodb": "data-custodian",
    "cassandra": "data-custodian",
    "neo4j": "data-custodian",
    "all": "data-custodian",
    "hadoop": "hadoop-ranger",
}
AGENT_NAME = {
    "data-custodian": "Data Custodian",
    "hadoop-ranger": "Hadoop Ranger",
    "etl-guardian": "ETL Guardian",
    "lakehouse-ops": "Lakehouse Ops",
}

# 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}
# count(*) over these connectors does a full scan (slow), so a background task
# refreshes the counts periodically and the endpoint always returns instantly.
_sync_cache: dict[str, Any] = {"counts": {k: None for k in SOURCE_COUNT_SQL}, "ts": 0.0}
_refresher_started = False
_SYNC_INTERVAL = 60.0


async def _sync_refresh_loop() -> None:
    while True:
        try:
            keys = list(SOURCE_COUNT_SQL.keys())
            results = await asyncio.gather(
                *[_trino_scalar(SOURCE_COUNT_SQL[k], deadline_s=90.0) for k in keys]
            )
            merged = dict(_sync_cache["counts"])
            for k, v in zip(keys, results):
                if v is not None:
                    merged[k] = v
            _sync_cache["counts"] = merged
            _sync_cache["ts"] = time.time()
        except Exception:
            pass
        await asyncio.sleep(_SYNC_INTERVAL)


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


def _feed(agent_id: str, message: str, level: str = "info") -> None:
    """Write an entry to the shared agent activity feed (Comms log)."""
    try:
        from main import add_feed  # lazy: main is fully loaded by request time
        add_feed(agent_id, message, level)
    except Exception:
        pass


async def _watch_run(source: str, dag_id: str, run_id: str, agent_id: str, rows: int | None) -> None:
    """Poll an Airflow run to completion and log the outcome to the feed."""
    name = AGENT_NAME.get(agent_id, agent_id)
    label = f"{rows} rows" if rows else "data"
    try:
        async with httpx.AsyncClient() as client:
            tok = await _airflow_token(client)
            headers = {"Authorization": f"Bearer {tok}"}
            for _ in range(180):  # up to ~15 min
                await asyncio.sleep(5)
                try:
                    r = await client.get(
                        f"{AIRFLOW_URL}/api/v2/dags/{dag_id}/dagRuns/{run_id}",
                        headers=headers, timeout=10,
                    )
                    state = r.json().get("state")
                except Exception:
                    continue
                if state == "success":
                    _feed(agent_id, f"[datagen] {name} generated {label} in {source} — complete, data flowing via CDC to Kafka/S3", "info")
                    return
                if state == "failed":
                    _feed(agent_id, f"[datagen] {name}: generation for {source} failed (see Airflow logs)", "err")
                    return
    except Exception:
        pass


async def _trino_scalar(sql: str, deadline_s: float = 8.0) -> int | None:
    """Run a scalar Trino query with a hard wall-clock deadline.

    count(*) over a large Cassandra table can scan for a long time; without a
    total deadline the result-paging loop would hang the endpoint. On timeout
    we abort the query and return None so the UI degrades gracefully.
    """
    end = time.time() + deadline_s
    try:
        async with httpx.AsyncClient(timeout=5.0) 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] = []
            if data.get("data"):
                rows += data["data"]
            nxt = data.get("nextUri")
            while nxt:
                if time.time() > end:
                    try:
                        await client.delete(nxt, timeout=3.0)
                    except Exception:
                        pass
                    return None
                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)
    agent_id = body.get("agent_id") or SOURCE_AGENT.get(source, "data-custodian")
    autonomous = bool(body.get("autonomous"))
    name = AGENT_NAME.get(agent_id, agent_id)
    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:
            _feed(agent_id, f"[datagen] {name}: could not start generation for {source} (Airflow {r.status_code})", "err")
            return JSONResponse({"ok": False, "error": f"Airflow {r.status_code}: {r.text[:300]}"}, status_code=200)
        j = r.json()
        run_id = j.get("dag_run_id")
        verb = "generating autonomously" if autonomous else "started generation of"
        rows_txt = f"{conf['rows']} rows" if conf.get("rows") else "data"
        _feed(agent_id, f"[datagen] {name} {verb} {rows_txt} in {source}", "info")
        if run_id:
            asyncio.create_task(_watch_run(source, dag_id, run_id, agent_id, conf.get("rows")))
        return JSONResponse({
            "ok": True,
            "source": source,
            "dag_id": dag_id,
            "run_id": run_id,
            "state": j.get("state"),
            "rows": conf.get("rows"),
            "agent_id": agent_id,
            "agent_name": name,
        })
    except Exception as exc:
        return JSONResponse({"ok": False, "error": str(exc)}, status_code=200)


@router.get("/agents")
async def agents() -> JSONResponse:
    """Which agent is responsible for generating each source."""
    out = {src: {"agent_id": aid, "agent_name": AGENT_NAME.get(aid, aid)}
           for src, aid in SOURCE_AGENT.items()}
    return JSONResponse({"ok": True, "agents": out})


@router.get("/activity")
async def activity(limit: int = Query(25)) -> JSONResponse:
    """Recent data-generation activity performed by agents (from the feed)."""
    try:
        from main import FeedEntry
        from db import SessionLocal
        from sqlalchemy import select
        with SessionLocal() as db:
            rows = db.execute(
                select(FeedEntry).order_by(FeedEntry.ts.desc()).limit(400)
            ).scalars().all()
        items = []
        for r in rows:
            if r.message and "[datagen]" in r.message:
                items.append({
                    "id": r.id,
                    "ts": r.ts.isoformat() if r.ts else None,
                    "agent_id": r.agent_id,
                    "agent_name": AGENT_NAME.get(r.agent_id, r.agent_id),
                    "message": r.message.replace("[datagen] ", ""),
                    "level": r.level,
                })
            if len(items) >= limit:
                break
        return JSONResponse({"ok": True, "activity": items})
    except Exception as exc:
        return JSONResponse({"ok": False, "error": str(exc), "activity": []}, 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), refresh: bool = Query(False)) -> JSONResponse:
    """Live row counts per source (via Trino) for the 'in sync' display.

    Counts are cached (TTL) and refreshed concurrently because count(*) over
    these connectors does a full scan. Returns immediately from cache when
    fresh; otherwise refreshes once (other concurrent callers reuse the cache).
    """
    global _refresher_started
    if not _refresher_started:
        _refresher_started = True
        asyncio.create_task(_sync_refresh_loop())
    counts = _sync_cache["counts"]
    age = round(time.time() - _sync_cache["ts"], 1) if _sync_cache["ts"] else None
    if source and source in counts:
        return JSONResponse({"ok": True, "counts": {source: counts[source]}, "age_s": age})
    return JSONResponse({"ok": True, "counts": counts, "age_s": age})