atc-agents/api/movements.py

"""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})