"""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} rijen" 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} genereerde {label} in {source} — klaar, data stroomt via CDC naar Kafka/S3", "info") return if state == "failed": _feed(agent_id, f"[datagen] {name}: generatie voor {source} is mislukt (zie 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}: kon generatie voor {source} niet starten (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 = "genereert zelf" if autonomous else "startte generatie:" rows_txt = f"{conf['rows']} rijen" 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})