feat(etl): Hadoop->Trino movement DAG + movements registry + autonomous ETL-agents

- DAG hadoop_to_trino (deployed to Airflow) + worker hadoop_to_trino.py (on the Hadoop master) move HDFS historical_sales -> iceberg.hadoop.historical_sales_hdfs. - api/movements.py: movement registry, Airflow trigger+watch, run tracking (state/duration/rows), endpoints /api/movements, /{id}/run, /runs. - agent_ops.py: ETL-agent loop autonomously triggers movements on an interval and logs each run; /api/agent-ops/etl/toggle + etl status.
2026-06-27 01:59:35 +02:00
parent 921342442f
commit a724615a9a
4 changed files with 257 additions and 3 deletions
@@ -4,7 +4,7 @@ WORKDIR /app
 RUN apt-get update && apt-get install -y --no-install-recommends curl && rm -rf /var/lib/apt/lists/*
 COPY requirements.txt .
 RUN pip install --no-cache-dir -r requirements.txt
-COPY main.py lab_context.py agent_terminal.py workload.py node_registry.py node_ops.py topology_views.py supervisor.py approval_service.py db.py dockhand_envs.py presentation.py database_inventory.py presentation_upload.py presentation_static.py storage_s3.py elasticsearch_api.py sql_console.py hdfs_api.py ssh_terminal.py pipeline_ops.py hadoop_analytics.py agent_ops.py cdc_consumer.py hive_bench_seed.json .
+COPY main.py lab_context.py agent_terminal.py workload.py node_registry.py node_ops.py topology_views.py supervisor.py approval_service.py db.py dockhand_envs.py presentation.py database_inventory.py presentation_upload.py presentation_static.py storage_s3.py elasticsearch_api.py sql_console.py hdfs_api.py ssh_terminal.py pipeline_ops.py hadoop_analytics.py agent_ops.py cdc_consumer.py movements.py hive_bench_seed.json .
 RUN mkdir -p /data
 ENV DATABASE_URL=sqlite:////data/atc-agents.db
 EXPOSE 3201
@@ -49,6 +49,10 @@ DML_AGENT = "data-custodian"

 _INTERVAL = float(os.getenv("AGENT_DML_INTERVAL_SECONDS", "45"))
 _MAX_ROWS = int(os.getenv("AGENT_DML_MAX_ROWS", "5"))
+# ETL-agent: autonomously trigger data movements on an interval.
+_ETL_INTERVAL = float(os.getenv("ETL_AGENT_INTERVAL_SECONDS", "300"))
+# Movements the ETL-agent cycles through autonomously (must exist in movements.py).
+_ETL_ROTATION = [m for m in os.getenv("ETL_AGENT_MOVEMENTS", "hadoop_to_trino,gen_postgres,gen_mysql,gen_mongodb").split(",") if m]
 # How many agent-created rows we keep around per source before favouring DELETE.
 _POOL_CAP = int(os.getenv("AGENT_DML_POOL_CAP", "400"))

@@ -64,6 +68,16 @@ _state: dict[str, Any] = {
    "by_op": {"insert": 0, "update": 0, "delete": 0},
 }

+_etl_state: dict[str, Any] = {
+    "enabled": os.getenv("ETL_AGENT_ENABLED", "1") not in ("0", "false", "False", ""),
+    "interval": _ETL_INTERVAL,
+    "rotation": _ETL_ROTATION,
+    "idx": 0,
+    "runs_total": 0,
+    "last": None,
+    "started": False,
+}
+
 # ── Realistic value domains ─────────────────────────────────────────────────
 REGIONS = ["NA", "EU", "EMEA", "APAC", "LATAM"]
 CHANNELS = ["online", "retail", "partner", "wholesale", "direct"]
@@ -291,11 +305,49 @@ async def agent_dml_loop() -> None:
        await asyncio.sleep(max(5.0, float(_state["interval"])))


+async def etl_agent_loop() -> None:
+    """Background loop: ETL-agents autonomously trigger data movements and log
+    each run (rows, duration, status) to the feed."""
+    _etl_state["started"] = True
+    await asyncio.sleep(30)  # let the platform settle
+    while True:
+        try:
+            if _etl_state["enabled"] and _etl_state["rotation"]:
+                from movements import trigger_and_watch, MOVEMENT_BY_ID
+
+                mid = _etl_state["rotation"][_etl_state["idx"] % len(_etl_state["rotation"])]
+                _etl_state["idx"] += 1
+                if mid in MOVEMENT_BY_ID:
+                    result = await trigger_and_watch(mid, autonomous=True)
+                    _etl_state["runs_total"] += 1
+                    _etl_state["last"] = {"movement_id": mid, "state": result.get("state"),
+                                          "rows": result.get("rows"), "duration_s": result.get("duration_s"),
+                                          "ts": datetime.now(timezone.utc).isoformat()}
+        except Exception as exc:
+            _etl_state["last"] = {"error": str(exc), "ts": datetime.now(timezone.utc).isoformat()}
+        await asyncio.sleep(max(30.0, float(_etl_state["interval"])))
+
+
 # ── Endpoints ────────────────────────────────────────────────────────────────
@router.get("/status")
 async def status() -> JSONResponse:
    pools = {k: len(v) for k, v in _pools.items()}
-    return JSONResponse({"ok": True, "pools": pools, **_state})
+    return JSONResponse({"ok": True, "pools": pools, "etl": _etl_state, **_state})
+
+
+@router.post("/etl/toggle")
+async def etl_toggle(body: dict[str, Any] = Body(default={})) -> JSONResponse:
+    if "enabled" in body:
+        _etl_state["enabled"] = bool(body["enabled"])
+    else:
+        _etl_state["enabled"] = not _etl_state["enabled"]
+    if "interval" in body:
+        try:
+            _etl_state["interval"] = max(30.0, float(body["interval"]))
+        except (TypeError, ValueError):
+            pass
+    await _emit(f"[etl-agent] Autonomous ETL movements {'ENABLED' if _etl_state['enabled'] else 'PAUSED'} by operator", "warn")
+    return JSONResponse({"ok": True, "enabled": _etl_state["enabled"], "interval": _etl_state["interval"]})


@router.post("/toggle")
@@ -42,8 +42,9 @@ from pipeline_ops import router as pipeline_router
 from hadoop_analytics import router as hadoop_router
 from elasticsearch_api import router as elasticsearch_router
 from sql_console import router as sql_router
-from agent_ops import router as agent_ops_router, agent_dml_loop
+from agent_ops import router as agent_ops_router, agent_dml_loop, etl_agent_loop
 from cdc_consumer import router as cdc_router, cdc_consumer_loop
+from movements import router as movements_router
 from ssh_terminal import ssh_session
 from node_registry import NODE_IDS, NODE_AGENT, NODE_REGISTRY, is_node_id
 from node_ops import build_node_detail, probe_node, run_node_probe_task
@@ -715,11 +716,13 @@ async def lifespan(app: FastAPI):
    task = asyncio.create_task(heartbeat_loop())
    dml_task = asyncio.create_task(agent_dml_loop())
    cdc_task = asyncio.create_task(cdc_consumer_loop())
+    etl_task = asyncio.create_task(etl_agent_loop())
    add_feed("infra-sentinel", "ATC Command Center API online", "info")
    yield
    task.cancel()
    dml_task.cancel()
    cdc_task.cancel()
+    etl_task.cancel()
    if redis_client:
        await redis_client.close()

@@ -733,6 +736,7 @@ app.include_router(elasticsearch_router)
 app.include_router(sql_router)
 app.include_router(agent_ops_router)
 app.include_router(cdc_router)
+app.include_router(movements_router)
 app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
@@ -0,0 +1,198 @@
+"""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": "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"},
+]
+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 _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}"}
+    conf = {**(mv.get("default_conf") or {}), **(conf or {})}
+    agent = mv["agent"]
+    count_sql = mv.get("count_sql")
+    before = await _trino_scalar(count_sql) if count_sql else None
+
+    _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"})
+    verb = "autonomously triggered" if autonomous else "triggered"
+    _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)
+    except Exception as exc:
+        _last_runs[mid] = {**_last_runs[mid], "state": "failed", "error": str(exc)}
+        _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"
+    _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})