mo
212 lines
7.9 KiB
Python
212 lines
7.9 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 -> 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
|
|
|
|
|
|
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)
|
|
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), 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})
|