"""Data movement registry + orchestration for the Command Center. A "movement" is a named data flow step (generate into a source DB, move HDFS -> Iceberg via Trino, mask -> curated, ...). Each maps to an Airflow DAG. This module triggers movements via the Airflow REST API, watches the runs to completion, records the latest run (state, duration, rows) and logs to the agent feed. It is shared by the ETL-agents (autonomous triggering) and the Data Flow tab (live status + manual triggering). """ from __future__ import annotations import asyncio import os import time from datetime import datetime, timezone from typing import Any import httpx from fastapi import APIRouter, Body from fastapi.responses import JSONResponse router = APIRouter(prefix="/api/movements", tags=["movements"]) 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") # from/to refer to logical Data Flow node ids (see dataflow.py). MOVEMENTS: list[dict[str, Any]] = [ {"id": "gen_postgres", "label": "Generate → PostgreSQL", "kind": "generate", "dag_id": "gen_postgres", "agent": "data-custodian", "from": "generator", "to": "postgres", "default_conf": {"rows": 3000}}, {"id": "gen_mysql", "label": "Generate → MySQL", "kind": "generate", "dag_id": "gen_mysql", "agent": "data-custodian", "from": "generator", "to": "mysql", "default_conf": {"rows": 3000}}, {"id": "gen_mongodb", "label": "Generate → MongoDB", "kind": "generate", "dag_id": "gen_mongodb", "agent": "data-custodian", "from": "generator", "to": "mongodb", "default_conf": {"rows": 3000}}, {"id": "hdfs_to_kafka", "label": "HDFS → Kafka export", "kind": "stream", "dag_id": None, "agent": "hadoop-ranger", "from": "hdfs", "to": "kafka", "api": "/api/pipeline/streaming/hdfs/to-kafka", "default_conf": {"source": "trino", "table": "iceberg.hadoop.historical_sales_hdfs", "topic": "hdfs.historical.sales"}}, {"id": "hadoop_to_trino", "label": "HDFS → Iceberg (Trino)", "kind": "movement", "dag_id": "hadoop_to_trino", "agent": "hadoop-ranger", "from": "hdfs", "to": "iceberg_hadoop", "default_conf": {"mode": "refresh"}, "count_sql": "SELECT count(*) FROM iceberg.hadoop.historical_sales_hdfs"}, {"id": "mask_to_curated", "label": "Mask PII → Curated (Iceberg)", "kind": "mask", "dag_id": "mask_to_curated", "agent": "lakehouse-ops", "from": "sources", "to": "iceberg_curated", "default_conf": {}, "count_sql": "SELECT count(*) FROM iceberg.curated_masked.sales_orders_masked"}, {"id": "spark_to_s3", "label": "Spark → S3 curated", "kind": "movement", "dag_id": "mask_to_curated", "agent": "lakehouse-ops", "from": "spark", "to": "s3_cdc", "default_conf": {"target": "s3"}, "count_sql": "SELECT count(*) FROM iceberg.curated_masked.sales_orders_masked"}, {"id": "spark_to_curated", "label": "Spark → Iceberg curated", "kind": "movement", "dag_id": "mask_to_curated", "agent": "lakehouse-ops", "from": "spark", "to": "iceberg_curated", "default_conf": {}, "count_sql": "SELECT count(*) FROM iceberg.curated_masked.sales_orders_masked"}, ] MOVEMENT_BY_ID = {m["id"]: m for m in MOVEMENTS} # latest run info per movement id _last_runs: dict[str, dict[str, Any]] = {} _token_cache: dict[str, Any] = {"token": None, "exp": 0.0} def _feed(agent_id: str, message: str, level: str = "info") -> None: try: from main import add_feed add_feed(agent_id, message, level) except Exception: pass async def _publish(event: dict[str, Any]) -> None: try: from main import publish_event await publish_event(event) except Exception: pass async def _term(agent_id: str, text: str, level: str = "info", phase: str = "etl") -> None: """Stream a movement step to the owning agent's terminal.""" try: from agent_terminal import terminal_log await terminal_log(agent_id, text, level=level, phase=phase) except Exception: pass 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 return tok async def _trino_scalar(sql: str, deadline_s: float = 15.0) -> int | None: end = time.time() + deadline_s try: async with httpx.AsyncClient(timeout=6.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] = data.get("data") or [] nxt = data.get("nextUri") while nxt: if time.time() > end: try: await client.delete(nxt, timeout=3.0) except Exception: pass return None d = (await client.get(nxt)).json() rows += d.get("data") or [] if d.get("error"): return None nxt = d.get("nextUri") if rows: return int(rows[0][0]) except Exception: return None return None async def trigger_and_watch(mid: str, conf: dict[str, Any] | None = None, *, autonomous: bool = False) -> dict[str, Any]: mv = MOVEMENT_BY_ID.get(mid) if not mv: return {"ok": False, "error": f"unknown movement {mid}"} agent = mv["agent"] if mv.get("api"): try: payload = {**(mv.get("default_conf") or {}), **(conf or {})} await _term(agent, f"$ POST {mv['api']} # {mv['label']}", level="cmd") await _term(agent, f" payload={payload}", level="cmd") t0 = time.time() async with httpx.AsyncClient(timeout=120.0) as client: r = await client.post(f"http://127.0.0.1:8000{mv['api']}", json=payload) dur = round(time.time() - t0, 1) body = r.json() if r.headers.get("content-type", "").startswith("application/json") else {} state = "success" if r.status_code < 400 and body.get("ok", True) else "failed" rows = body.get("rows_sent") or body.get("rows") run = { "movement_id": mid, "state": state, "duration_s": dur, "rows": rows, "ended_at": datetime.now(timezone.utc).isoformat(), "conf": payload, } _last_runs[mid] = run await _publish({"type": "movement", **run}) lvl = "info" if state == "success" else "err" await _term(agent, f" ← {state} · {rows if rows is not None else '?'} rows in {dur}s", level="ok" if state == "success" else "err") _feed(agent, f"[etl] {mv['label']}: {state} in {dur}s", lvl) return {"ok": state == "success", **run} except Exception as exc: _last_runs[mid] = {**_last_runs.get(mid, {}), "state": "failed", "error": str(exc)} await _term(agent, f" ✗ error {str(exc)[:140]}", level="err") _feed(agent, f"[etl] {mv['label']}: error {str(exc)[:120]}", "err") return {"ok": False, "error": str(exc)} conf = {**(mv.get("default_conf") or {}), **(conf or {})} count_sql = mv.get("count_sql") verb = "autonomously triggered" if autonomous else "triggered" await _term(agent, f"$ airflow dags trigger {mv['dag_id']} # {mv['label']} ({verb})", level="cmd") await _term(agent, f" conf={conf}", level="cmd") before = await _trino_scalar(count_sql) if count_sql else None if count_sql: await _term(agent, f" trino: {count_sql} → {before if before is not None else '?'} rows (before)", level="info") _last_runs[mid] = {**_last_runs.get(mid, {}), "state": "running", "started_at": datetime.now(timezone.utc).isoformat()} await _publish({"type": "movement", "movement_id": mid, "state": "running"}) _feed(agent, f"[etl] {mv['label']}: {verb} (conf={conf})", "info") try: async with httpx.AsyncClient() as client: tok = await _airflow_token(client) h = {"Authorization": f"Bearer {tok}"} r = await client.post(f"{AIRFLOW_URL}/api/v2/dags/{mv['dag_id']}/dagRuns", headers=h, json={"logical_date": None, "conf": conf}, timeout=20) if r.status_code >= 400: _last_runs[mid] = {**_last_runs[mid], "state": "failed", "error": f"airflow {r.status_code}"} _feed(agent, f"[etl] {mv['label']}: could not start (Airflow {r.status_code})", "err") return {"ok": False, "error": f"airflow {r.status_code}: {r.text[:200]}"} run_id = r.json().get("dag_run_id") t0 = time.time() state = "running" for _ in range(120): # up to ~10 min await asyncio.sleep(5) rr = await client.get(f"{AIRFLOW_URL}/api/v2/dags/{mv['dag_id']}/dagRuns/{run_id}", headers=h, timeout=10) state = rr.json().get("state") if state in ("success", "failed"): break dur = round(time.time() - t0, 1) await _term(agent, f" ← Airflow run {run_id} finished state={state} in {dur}s", level="ok" if state == "success" else "err") except Exception as exc: _last_runs[mid] = {**_last_runs[mid], "state": "failed", "error": str(exc)} await _term(agent, f" ✗ error {str(exc)[:140]}", level="err") _feed(agent, f"[etl] {mv['label']}: error {str(exc)[:120]}", "err") return {"ok": False, "error": str(exc)} after = await _trino_scalar(count_sql) if count_sql else None rows = None if before is not None and after is not None: rows = max(0, after - before) if conf.get("mode") != "refresh" else after elif conf.get("rows"): rows = conf.get("rows") run = { "movement_id": mid, "state": state, "duration_s": dur, "rows": rows, "ended_at": datetime.now(timezone.utc).isoformat(), "run_id": run_id, "conf": conf, } _last_runs[mid] = run await _publish({"type": "movement", **run}) rtxt = f"{rows} rows" if rows is not None else "data" lvl = "info" if state == "success" else "err" if count_sql: await _term(agent, f" trino: count {before if before is not None else '?'} → {after if after is not None else '?'} ({rtxt})", level="info") await _term(agent, f" ← {mv['label']}: {state} — {rtxt} in {dur}s", level="ok" if state == "success" else "err") _feed(agent, f"[etl] {mv['label']}: {state} — {rtxt} in {dur}s", lvl) return {"ok": state == "success", **run} def last_runs() -> dict[str, dict[str, Any]]: return dict(_last_runs) # ── Endpoints ──────────────────────────────────────────────────────────────── @router.get("") async def list_movements() -> JSONResponse: out = [] for m in MOVEMENTS: out.append({**{k: m[k] for k in ("id", "label", "kind", "dag_id", "agent", "from", "to")}, "last_run": _last_runs.get(m["id"])}) return JSONResponse({"ok": True, "movements": out}) @router.post("/{mid}/run") async def run_movement(mid: str, body: dict[str, Any] = Body(default={})) -> JSONResponse: if mid not in MOVEMENT_BY_ID: return JSONResponse({"ok": False, "error": f"unknown movement {mid}"}, status_code=400) conf = body.get("conf") if isinstance(body, dict) else None # Run in the background so the request returns immediately; status via WS/feed. asyncio.create_task(trigger_and_watch(mid, conf)) return JSONResponse({"ok": True, "movement_id": mid, "status": "triggered"}) @router.get("/runs") async def movement_runs() -> JSONResponse: return JSONResponse({"ok": True, "last_runs": _last_runs})