atc-agents/api/dataflow.py

"""Data Flow graph for the Command Center "Data Flow" tab.

Assembles the live landscape as a node-link graph (generators -> source DBs ->
CDC/Kafka -> sinks; HDFS -> Iceberg via Trino; sources -> masked curated layer)
and overlays live status: movement run states, CDC volume, Trino row counts, and
a PII overlay (which nodes hold PII and whether it is masked).

This is the single LIVE source of truth for the Data Flow tab. The structure is
defined here (movement-driven) and is OpenMetadata-ready: when lineage is
available it can enrich/replace the static edges.
"""

from __future__ import annotations

import os
import time
from pathlib import Path
from typing import Any

import httpx
from fastapi import APIRouter, Body
from fastapi.responses import JSONResponse

router = APIRouter(prefix="/api/dataflow", tags=["dataflow"])

TRINO_URL = os.getenv("TRINO_URL", "http://10.0.21.50:8089").rstrip("/")
TRINO_USER = os.getenv("TRINO_USER", "mo")

# Static landscape (x,y in 0..100). kind drives UI styling.
# Laid out as clean left→right pipeline stages so lineage reads in order:
#   producers (col 0) → sources (col 1) → CDC (col 2) → storage/lakehouse (col 3)
#   → query engine (col 4); governance (OpenMetadata) sits centred at the bottom.
NODES: list[dict[str, Any]] = [
    # col 0 — producers / origin
    {"id": "generator", "label": "Data Generator", "sub": "Airflow DAGs", "kind": "generator", "x": 9, "y": 34},
    {"id": "hdfs", "label": "Hadoop HDFS", "sub": "historical_sales", "kind": "hadoop", "x": 9, "y": 80},
    # col 1 — source databases
    {"id": "postgres", "label": "PostgreSQL", "sub": "sales_orders", "kind": "source", "x": 28, "y": 12},
    {"id": "mysql", "label": "MySQL", "sub": "employee_events", "kind": "source", "x": 28, "y": 28},
    {"id": "mongodb", "label": "MongoDB", "sub": "events", "kind": "source", "x": 28, "y": 44},
    {"id": "cassandra", "label": "Cassandra", "sub": "device_metrics", "kind": "source", "x": 28, "y": 60},
    {"id": "neo4j", "label": "Neo4j", "sub": "Product · Supplier graph", "kind": "source", "x": 28, "y": 76},
    # col 2 — change data capture
    {"id": "kafka", "label": "Kafka · Debezium", "sub": "CDC topics", "kind": "stream", "x": 47, "y": 24},
    {"id": "spark", "label": "Apache Spark", "sub": "Streaming · batch", "kind": "compute", "x": 62, "y": 38,
     "url": "/spark-ui/"},
    # col 3 — storage / lakehouse
    {"id": "s3_cdc", "label": "S3 CDC Archive", "sub": "object store", "kind": "sink", "x": 67, "y": 13},
    {"id": "iceberg_curated", "label": "Iceberg · curated_masked", "sub": "masked PII", "kind": "lakehouse", "x": 67, "y": 45},
    {"id": "iceberg_hadoop", "label": "Iceberg · hadoop", "sub": "historical_sales_hdfs", "kind": "lakehouse", "x": 67, "y": 80},
    # col 4 — query engine
    {"id": "trino", "label": "Trino", "sub": "query engine", "kind": "engine", "x": 88, "y": 45},
    # governance — bottom centre
    {"id": "openmetadata", "label": "OpenMetadata", "sub": "catalog · lineage · PII", "kind": "governance",
     "x": 47, "y": 93, "url": "http://10.0.21.47:8585"},
    # AI serving lane — how governed data reaches the LLM & the Command Center chat
    {"id": "chromadb", "label": "ChromaDB", "sub": "vectors · embeddings", "kind": "vector", "x": 60, "y": 65},
    {"id": "rag", "label": "RAG · LangChain", "sub": "retrieve · augment · agent", "kind": "rag", "x": 74, "y": 65},
    {"id": "vllm", "label": "vLLM Gateway", "sub": "Llama3-70B · GPT-4o", "kind": "llm", "x": 88, "y": 68},
    {"id": "chat", "label": "Knowledge Chat", "sub": "Command Center", "kind": "chat", "x": 92, "y": 90},
]

# Edges. movement_id (optional) links to movements.py so the edge is triggerable.
EDGES: list[dict[str, Any]] = [
    {"from": "generator", "to": "postgres", "kind": "generate", "movement_id": "gen_postgres"},
    {"from": "generator", "to": "mysql", "kind": "generate", "movement_id": "gen_mysql"},
    {"from": "generator", "to": "mongodb", "kind": "generate", "movement_id": "gen_mongodb"},
    {"from": "generator", "to": "cassandra", "kind": "generate"},
    {"from": "generator", "to": "neo4j", "kind": "generate"},
    {"from": "postgres", "to": "kafka", "kind": "cdc"},
    {"from": "mysql", "to": "kafka", "kind": "cdc"},
    {"from": "mongodb", "to": "kafka", "kind": "cdc"},
    {"from": "cassandra", "to": "kafka", "kind": "cdc"},
    {"from": "neo4j", "to": "kafka", "kind": "cdc"},
    {"from": "kafka", "to": "spark", "kind": "stream"},
    {"from": "spark", "to": "iceberg_curated", "kind": "movement", "movement_id": "spark_to_curated"},
    {"from": "spark", "to": "s3_cdc", "kind": "movement", "movement_id": "spark_to_s3"},
    {"from": "kafka", "to": "s3_cdc", "kind": "archive"},
    {"from": "postgres", "to": "hdfs", "kind": "archive", "offload": True},
    {"from": "mysql", "to": "hdfs", "kind": "archive", "offload": True},
    {"from": "hdfs", "to": "kafka", "kind": "stream", "movement_id": "hdfs_to_kafka"},
    {"from": "hdfs", "to": "iceberg_hadoop", "kind": "movement", "movement_id": "hadoop_to_trino"},
    {"from": "postgres", "to": "iceberg_curated", "kind": "mask", "movement_id": "mask_to_curated"},
    {"from": "mysql", "to": "iceberg_curated", "kind": "mask", "movement_id": "mask_to_curated"},
    {"from": "mongodb", "to": "iceberg_curated", "kind": "mask", "movement_id": "mask_to_curated"},
    {"from": "iceberg_hadoop", "to": "trino", "kind": "query"},
    {"from": "iceberg_curated", "to": "trino", "kind": "query"},
    {"from": "kafka", "to": "trino", "kind": "query"},
    {"from": "postgres", "to": "openmetadata", "kind": "catalog"},
    {"from": "mysql", "to": "openmetadata", "kind": "catalog"},
    {"from": "mongodb", "to": "openmetadata", "kind": "catalog"},
    {"from": "cassandra", "to": "openmetadata", "kind": "catalog"},
    {"from": "neo4j", "to": "openmetadata", "kind": "catalog"},
    {"from": "trino", "to": "openmetadata", "kind": "catalog"},
    # AI serving lane: governed business data + catalog + vectors → RAG → vLLM → chat
    {"from": "trino", "to": "rag", "kind": "context"},
    {"from": "openmetadata", "to": "rag", "kind": "context"},
    {"from": "iceberg_curated", "to": "rag", "kind": "context"},
    {"from": "chromadb", "to": "rag", "kind": "retrieve"},
    {"from": "rag", "to": "vllm", "kind": "prompt"},
    {"from": "vllm", "to": "chat", "kind": "answer"},
    {"from": "rag", "to": "chat", "kind": "answer"},
]

_cache: dict[str, Any] = {"ts": 0.0, "data": None}
_TTL = 12.0
_count_cache: dict[str, Any] = {"ts": 0.0, "val": None}


def _trino_scalar(sql: str, deadline_s: float = 8.0) -> int | None:
    end = time.time() + deadline_s
    try:
        with httpx.Client(timeout=5.0) as client:
            d = client.post(f"{TRINO_URL}/v1/statement", content=sql.encode(), headers={"X-Trino-User": TRINO_USER}).json()
            rows: list[Any] = d.get("data") or []
            nxt = d.get("nextUri")
            while nxt:
                if time.time() > end:
                    try:
                        client.delete(nxt, timeout=3.0)
                    except Exception:
                        pass
                    return None
                dd = client.get(nxt).json()
                rows += dd.get("data") or []
                if dd.get("error"):
                    return None
                nxt = dd.get("nextUri")
            return int(rows[0][0]) if rows else None
    except Exception:
        return None


RAG_URL = os.getenv("RAG_URL", "http://rag-api:5020").rstrip("/")
_rag_cache: dict[str, Any] = {"ts": 0.0, "val": None}


def _rag_info() -> dict[str, Any]:
    now = time.time()
    if _rag_cache["val"] is not None and now - _rag_cache["ts"] < 60:
        return _rag_cache["val"]
    info: dict[str, Any] = {}
    try:
        with httpx.Client(timeout=2.0) as client:
            info = client.get(f"{RAG_URL}/config").json() or {}
    except Exception:
        info = {}
    _rag_cache["val"] = info
    _rag_cache["ts"] = now
    return info


def _iceberg_hadoop_count() -> int | None:
    now = time.time()
    if _count_cache["val"] is not None and now - _count_cache["ts"] < 60:
        return _count_cache["val"]
    v = _trino_scalar("SELECT count(*) FROM iceberg.hadoop.historical_sales_hdfs")
    if v is not None:
        _count_cache["val"] = v
        _count_cache["ts"] = now
    return v


async def _build() -> dict[str, Any]:
    # Live signals
    try:
        from movements import last_runs
        runs = last_runs()
    except Exception:
        runs = {}
    try:
        from cdc_consumer import snapshot as cdc_snapshot
        cdc = cdc_snapshot(15)
    except Exception:
        cdc = {"connected": False, "consumed": 0, "by_source": {}, "window_total": 0}
    try:
        from trino_federated import generator_active
        gen_active = generator_active()
    except Exception:
        gen_active = False
    try:
        from storage_s3 import archive_active
        arch_active = archive_active()
    except Exception:
        arch_active = False
    try:
        from pii_catalog import get_pii
        pii = get_pii()
    except Exception:
        pii = {"datasets": []}
    try:
        from streaming_ops import build_streaming_status
        streaming = await build_streaming_status()
    except Exception:
        streaming = {}
    spark = streaming.get("spark") or {}
    kafka = streaming.get("kafka") or {}
    edge_live = streaming.get("edges") or {}
    pii_by_node = {d["node_id"]: d for d in pii.get("datasets", [])}

    nodes = []
    for n in NODES:
        node = dict(n)
        node["level"] = "ok"
        metric = None
        if n["id"] == "openmetadata":
            metric = f"{pii.get('summary', {}).get('pii_columns', 0)} PII cols cataloged"
        elif n["id"] == "kafka":
            topics = len(kafka.get("topics") or [])
            conn_n = len(kafka.get("connectors") or [])
            metric = f"{cdc.get('window_total', 0)} chg/15m · {topics} topics · {conn_n} connectors"
            node["level"] = "ok" if cdc.get("connected") and kafka.get("ui_ok") else "warn"
        elif n["id"] == "spark":
            apps = len(spark.get("active_apps") or [])
            cores = spark.get("cores") or 0
            used = spark.get("cores_used") or 0
            metric = f"{spark.get('alive_workers', 0)} workers · {used}/{cores} cores · {apps} apps"
            node["level"] = "ok" if spark.get("ui_ok") and (spark.get("status") or "").upper() == "ALIVE" else "warn"
        elif n["id"] in ("postgres", "mysql", "mongodb", "cassandra", "neo4j"):
            metric = f"{cdc.get('by_source', {}).get(n['id'], 0)} CDC/15m"
        elif n["id"] == "iceberg_hadoop":
            c = _iceberg_hadoop_count()
            metric = f"{c:,} rows" if c is not None else "iceberg table"
        elif n["id"] == "generator":
            metric = "Airflow gen DAGs"
        elif n["id"] in ("vllm", "rag", "chromadb"):
            info = _rag_info()
            model = info.get("llm_model") or "gpt-4o"
            embed = (info.get("embed_model") or "all-MiniLM-L6-v2").split("/")[-1]
            if n["id"] == "vllm":
                metric = f"{model} · OpenAI-compat"
            elif n["id"] == "rag":
                metric = f"LangChain · {embed}"
            else:
                metric = f"embeddings · {embed}"
            node["level"] = "ok" if info else "warn"
        elif n["id"] == "chat":
            metric = "RAG chat · agent mode"
        # PII overlay
        p = pii_by_node.get(n["id"])
        if p:
            node["pii"] = {
                "key": p["key"],
                "has_pii": p["has_pii"], "pii_count": p["pii_count"], "all_masked": p["all_masked"],
                "masked_layer": p.get("masked_layer", False),
                "categories": sorted({c["category"] for c in p["pii_columns"]}),
                "columns": p["pii_columns"],
            }
        node["metric"] = metric
        nodes.append(node)

    try:
        from streaming_ops import flow_mode
        _flow = flow_mode()
    except Exception:
        _flow = "running"

    edges = []
    for e in EDGES:
        edge = dict(e)
        mid = e.get("movement_id")
        if mid and mid in runs:
            lr = runs[mid]
            edge["state"] = lr.get("state")
            edge["last_rows"] = lr.get("rows")
            edge["last_duration_s"] = lr.get("duration_s")
            edge["active"] = lr.get("state") == "running"
        if e["kind"] == "generate":
            # The live generator (continuous loop or the 'Generate data' burst)
            # writes into these four sources; pulse the edge while it is active.
            if e["to"] in ("postgres", "mysql", "mongodb", "cassandra"):
                edge["active"] = gen_active or bool(edge.get("active"))
        elif e["kind"] == "cdc":
            edge["active"] = cdc.get("by_source", {}).get(e["from"], 0) > 0
        elif e.get("from") == "hdfs" and e.get("to") == "kafka":
            edge["active"] = bool(edge_live.get("hdfs→kafka"))
        elif e.get("from") == "kafka" and e.get("to") == "spark":
            edge["active"] = bool(edge_live.get("kafka→spark"))
        elif e.get("from") == "spark" and e.get("to") == "iceberg_curated":
            edge["active"] = bool(edge_live.get("spark→iceberg")) or edge.get("active")
        elif e.get("from") == "spark" and e.get("to") == "s3_cdc":
            edge["active"] = bool(edge_live.get("spark→s3")) or edge.get("active") or arch_active
        elif e["kind"] == "archive" and e.get("from") == "kafka" and e.get("to") == "s3_cdc":
            # Kafka → S3 CDC archive pulses while the pipeline lands objects in S3
            edge["active"] = arch_active
        elif e["kind"] in ("context", "retrieve", "prompt", "answer"):
            # AI serving lane pulses while governed data is being served to the LLM
            edge["active"] = bool(_rag_info())
        if e.get("offload"):
            try:
                from agent_ops import custodian_recent
                edge["active"] = custodian_recent()
            except Exception:
                pass
        if _flow != "running":
            edge["active"] = False
        edges.append(edge)

    return {
        "ok": True,
        "nodes": nodes,
        "edges": edges,
        "pii_summary": pii.get("summary", {}),
        "cdc": {"connected": cdc.get("connected"), "consumed": cdc.get("consumed"), "window_total": cdc.get("window_total")},
        "streaming": streaming,
        "flow": _flow,
        "ts": time.time(),
    }


@router.get("")
async def get_dataflow(refresh: bool = False) -> JSONResponse:
    now = time.time()
    if not refresh and _cache["data"] and now - _cache["ts"] < _TTL:
        return JSONResponse(_cache["data"])
    data = await _build()
    _cache["data"] = data
    _cache["ts"] = now
    return JSONResponse(data)


GEN_SCRIPTS_DIR = os.getenv("GEN_SCRIPTS_DIR", "/app/gen_scripts")

# Which generation scripts to surface, in display order. `file` is relative to
# GEN_SCRIPTS_DIR (the mounted infra/airflow dir); `live` slices the running
# streamer source straight out of trino_federated.py so it is always in sync.
_SCRIPT_SPECS: list[dict[str, Any]] = [
    {"id": "live", "title": "Live streaming generator", "engine": "Command Center API",
     "desc": "Runs inside this API. While the Live dashboard is open it streams randomly-sized bursts of rows into PostgreSQL, MySQL, MongoDB & Cassandra every few seconds (and powers the manual 'Generate data' button). CDC propagates everything downstream.",
     "live": True},
    {"id": "dag", "title": "Airflow per-source DAGs", "engine": "Apache Airflow",
     "desc": "One triggerable DAG per database. Each passes a row count via dag_run conf and shells out to the matching generator script below.",
     "file": "per_source_gen_dags.py"},
    {"id": "postgres", "title": "PostgreSQL — sales orders", "engine": "Faker → psycopg2",
     "desc": "Generates realistic customers, products, regions, channels & amounts (incl. the PII columns that the masking layer later protects).",
     "file": "scripts/generate_postgres_sales_data.py"},
    {"id": "mysql", "title": "MySQL — employee events", "engine": "Faker → PyMySQL",
     "desc": "HR lifecycle events (hire, promotion, salary change, …) with employee PII.",
     "file": "scripts/generate_mysql_employee_data.py"},
    {"id": "mongodb", "title": "MongoDB — supply events", "engine": "Faker → PyMongo",
     "desc": "Schemaless supply-chain events with free-form payloads.",
     "file": "scripts/generate_mongodb_events_data.py"},
    {"id": "cassandra", "title": "Cassandra — device telemetry", "engine": "Faker → cassandra-driver",
     "desc": "High-volume IoT device metrics (temperature, voltage, …) on a time-series schema.",
     "file": "scripts/generate_cassandra_telemetry_data.py"},
    {"id": "neo4j", "title": "Neo4j — product & supplier graph", "engine": "Faker → neo4j driver",
     "desc": "Product/supplier nodes and relationships for the graph database.",
     "file": "scripts/generate_neo4j_graph_data.py"},
]


def _live_generator_source() -> str:
    try:
        text = Path("/app/trino_federated.py").read_text(encoding="utf-8")
    except Exception:
        return "# live generator source unavailable"
    start = text.find("# Continuous live generator")
    end = text.find('@router.get("/live")', start if start >= 0 else 0)
    if start >= 0 and end > start:
        return text[start:end].rstrip()
    return "# live generator source unavailable"


_scripts_cache: dict[str, Any] = {"ts": 0.0, "data": None}


@router.get("/scripts")
async def get_scripts() -> JSONResponse:
    now = time.time()
    if _scripts_cache["data"] and now - _scripts_cache["ts"] < 30:
        return JSONResponse(_scripts_cache["data"])
    base = Path(GEN_SCRIPTS_DIR)
    scripts = []
    for spec in _SCRIPT_SPECS:
        src = ""
        if spec.get("live"):
            src = _live_generator_source()
        else:
            p = base / spec["file"]
            try:
                src = p.read_text(encoding="utf-8")
            except Exception as exc:
                src = f"# source unavailable ({exc})"
        scripts.append({
            "id": spec["id"], "title": spec["title"], "engine": spec["engine"],
            "desc": spec["desc"], "filename": spec.get("file", "trino_federated.py"),
            "language": "python", "lines": src.count("\n") + 1, "source": src,
        })
    data = {"ok": True, "scripts": scripts}
    _scripts_cache["data"] = data
    _scripts_cache["ts"] = now
    return JSONResponse(data)


@router.post("/{movement_id}/run")
async def run_dataflow_movement(movement_id: str, body: dict[str, Any] = Body(default={})) -> JSONResponse:
    try:
        from movements import MOVEMENT_BY_ID, trigger_and_watch
        import asyncio
    except Exception as exc:
        return JSONResponse({"ok": False, "error": str(exc)}, status_code=500)
    if movement_id not in MOVEMENT_BY_ID:
        return JSONResponse({"ok": False, "error": f"unknown movement {movement_id}"}, status_code=400)
    conf = body.get("conf") if isinstance(body, dict) else None
    asyncio.create_task(trigger_and_watch(movement_id, conf))
    return JSONResponse({"ok": True, "movement_id": movement_id, "status": "triggered"})