diff --git a/api/Dockerfile b/api/Dockerfile index 07299c7..583dc8c 100644 --- a/api/Dockerfile +++ b/api/Dockerfile @@ -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 diff --git a/api/agent_ops.py b/api/agent_ops.py index 2a35a50..3c5e52e 100644 --- a/api/agent_ops.py +++ b/api/agent_ops.py @@ -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") diff --git a/api/main.py b/api/main.py index 20003de..71613bd 100644 --- a/api/main.py +++ b/api/main.py @@ -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=["*"], diff --git a/api/movements.py b/api/movements.py new file mode 100644 index 0000000..93f25b4 --- /dev/null +++ b/api/movements.py @@ -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})