mo
315 lines
12 KiB
Python
315 lines
12 KiB
Python
"""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})
|