atc-agents/api/lineage.py

"""Data lineage graph for the Command Center (Disease #6 — Traceability).

Models the real end-to-end pipeline as an explicit, staged node/edge graph:

  source DB → Debezium/Kafka (CDC) → Spark / Kafka Connect → S3 Parquet lake
            → Iceberg lakehouse (curated_masked / hadoop) → Trino federation
            → serving (Elasticsearch, RAG/ChromaDB, vLLM, Command Center chat)

Table-backed nodes are enriched with live Trino row counts. For the sales path
we expose column-level lineage with PII / masking status (reusing pii_catalog),
so an operator can trace any field from source to the masked curated layer.
OpenMetadata table lineage is layered in best-effort where the FQN resolves.

Endpoints:
  GET /api/lineage/graph?dataset=<key>   -> nodes + edges (+ column links)
  GET /api/lineage/datasets              -> selectable datasets
"""

from __future__ import annotations

import os
import time
from typing import Any

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

from lake_meta import DATASETS, DATASET_BY_KEY, trino_scalar

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

OPENMETADATA_URL = os.getenv("OPENMETADATA_URL", "").rstrip("/")
OPENMETADATA_TOKEN = os.getenv("OPENMETADATA_TOKEN", "")

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

# Which source datasets flow through the CDC → lakehouse pipeline.
_SOURCE_KEYS = ["orders", "hr", "supply", "telemetry"]


def _is_active() -> tuple[bool, bool]:
    gen = arch = False
    try:
        from trino_federated import generator_active
        gen = bool(generator_active())
    except Exception:
        pass
    try:
        from storage_s3 import archive_active
        arch = bool(archive_active())
    except Exception:
        pass
    return gen, arch


def _obs_counts() -> dict[str, int | None]:
    """Reuse the row counts the observability monitor already polls in-memory,
    so the lineage graph never issues its own (potentially slow) count(*)."""
    try:
        from observability import _state, _lock
        with _lock:
            return {k: v.get("rows") for k, v in _state["datasets"].items()}
    except Exception:
        return {}


def _row_count(ds_key: str, counts: dict[str, int | None]) -> int | None:
    return counts.get(ds_key)


def _pii_columns(ds_key: str) -> list[dict[str, Any]]:
    """Column-level detail with PII/masking flags from pii_catalog (best-effort).

    ds_key is a lake_meta key; map it to the pii_catalog key first."""
    pii_key = (DATASET_BY_KEY.get(ds_key, {}) or {}).get("pii_key", ds_key)
    try:
        from pii_catalog import get_pii
        data = get_pii()
        for d in data.get("datasets", []):
            if d.get("key") == pii_key or d.get("node_id") == pii_key:
                return [{"name": c["name"], "category": c.get("category"),
                         "masked": bool(c.get("masked")), "pii": True}
                        for c in d.get("pii_columns", [])]
    except Exception:
        pass
    return []


def _om_lineage(fqn: str) -> dict[str, int]:
    """OpenMetadata table lineage (best-effort, short timeout). Kept out of the
    hot path by default; only called when LINEAGE_OM_ENRICH=1 is set."""
    if not OPENMETADATA_URL or not fqn or os.getenv("LINEAGE_OM_ENRICH", "0") != "1":
        return {}
    url = f"{OPENMETADATA_URL}/api/v1/lineage/table/name/{fqn}?upstreamDepth=1&downstreamDepth=1"
    headers = {"Accept": "application/json"}
    if OPENMETADATA_TOKEN:
        headers["Authorization"] = f"Bearer {OPENMETADATA_TOKEN}"
    try:
        with httpx.Client(timeout=4.0, verify=False) as c:
            r = c.get(url, headers=headers)
            if r.status_code != 200:
                return {}
            j = r.json()
            return {"upstream": len(j.get("upstreamEdges") or []),
                    "downstream": len(j.get("downstreamEdges") or [])}
    except Exception:
        return {}


def _build() -> dict[str, Any]:
    gen_active, arch_active = _is_active()
    counts = _obs_counts()
    nodes: list[dict[str, Any]] = []
    edges: list[dict[str, Any]] = []
    seen: set[str] = set()

    def node(nid: str, label: str, ntype: str, stage: int, **meta: Any) -> None:
        if nid in seen:
            return
        seen.add(nid)
        nodes.append({"id": nid, "label": label, "type": ntype, "stage": stage, "meta": meta})

    def edge(src: str, dst: str, label: str, kind: str, active: bool = False) -> None:
        edges.append({"id": f"{src}->{dst}", "source": src, "target": dst,
                      "label": label, "kind": kind, "active": active})

    # Stage 5/6 shared serving nodes
    node("trino", "Trino", "engine", 5, engine="Trino", note="Federated SQL across every catalog")
    node("es", "Elasticsearch", "serving", 6, engine="Elasticsearch", note="Business + infra search indices")
    node("rag", "ChromaDB (RAG)", "serving", 6, engine="ChromaDB", note="Vector store / embeddings")
    node("vllm", "vLLM", "serving", 6, engine="vLLM", note="LLM inference (GPU)")
    node("cc", "Command Center", "serving", 6, engine="React", note="Knowledge Chat & dashboards")
    edge("rag", "vllm", "context", "serve", active=True)
    edge("vllm", "cc", "answers", "serve", active=True)
    edge("trino", "es", "index business data", "serve")
    edge("es", "cc", "search", "serve")
    edge("trino", "cc", "dashboards", "serve")

    for key in _SOURCE_KEYS:
        ds = DATASET_BY_KEY[key]
        src_id = f"src_{key}"
        topic_id = f"topic_{key}"
        lake_id = f"lake_{key}"
        rows = _row_count(key, counts)
        cols = _pii_columns(key)
        node(src_id, f"{ds['engine']}\n{ds['table']}", "source", 0,
             engine=ds["engine"], table=ds["fqtn"], rows=rows, domain=ds.get("domain"),
             columns=cols, ts_col=ds.get("ts_col"), key_col=ds.get("key_col"),
             om=_om_lineage(ds.get("om_fqn", "")))

        # CDC path (Debezium → Kafka topic)
        if ds.get("topic"):
            node(topic_id, f"Kafka\n{ds['topic']}", "stream", 1,
                 engine="Kafka", topic=ds["topic"], note="Debezium CDC topic")
            edge(src_id, topic_id, "Debezium CDC", "cdc", active=gen_active)
            node("connect", "Kafka Connect", "stream", 2, engine="Kafka Connect",
                 note="Debezium source + HDFS/S3 sinks")
            edge(topic_id, "connect", "consume", "cdc", active=gen_active)

        # ETL offload (direct source → S3 Parquet lake)
        node(lake_id, f"S3 lake\nlake/{key}", "storage", 3, engine="ObjectScale S3",
             path=f"s3://data/lake/{key}/", note="Parquet parts (ETL offload)")
        edge(src_id, lake_id, "ETL offload (Parquet)", "batch", active=arch_active)

    # Spark / Connect → curated + hadoop Iceberg tables
    node("spark", "Spark", "compute", 2, engine="Spark", note="Transform & mask → curated")
    edge("connect", "spark", "stream", "transform", active=gen_active)

    cur = DATASET_BY_KEY["curated"]
    had = DATASET_BY_KEY["hadoop"]
    node("iceberg_curated", f"Iceberg\n{cur['table']}", "lakehouse", 4, engine="Iceberg",
         table=cur["fqtn"], rows=_row_count("curated", counts), masked_layer=True,
         note="PII-masked curated layer", columns=_pii_columns("curated"),
         om=_om_lineage(cur.get("om_fqn", "")))
    node("iceberg_hadoop", f"Iceberg/HDFS\n{had['table']}", "lakehouse", 4, engine="Iceberg / HDFS",
         table=had["fqtn"], rows=_row_count("hadoop", counts), note="Historical sales on HDFS",
         om=_om_lineage(had.get("om_fqn", "")))

    edge("spark", "iceberg_curated", "mask + write", "transform", active=arch_active)
    edge("lake_orders", "iceberg_hadoop", "register external", "batch", active=arch_active)
    edge("src_orders", "iceberg_curated", "curate (masked)", "transform")

    # Lakehouse → Trino
    for nid in ("iceberg_curated", "iceberg_hadoop"):
        edge(nid, "trino", "query", "serve")
    for key in _SOURCE_KEYS:
        edge(f"src_{key}", "trino", "federate", "serve")

    # Lake → RAG (documents/parquet feeding embeddings is conceptual)
    edge("iceberg_curated", "rag", "embed (masked-safe)", "serve")

    # Column-level links: source sales_orders → curated masked
    column_links: list[dict[str, Any]] = []
    src_cols = {c["name"] for c in _pii_columns("orders")}
    cur_cols = {c["name"] for c in _pii_columns("curated")}
    for cn in sorted(src_cols & cur_cols):
        masked = any(c["name"] == cn and c["masked"] for c in _pii_columns("curated"))
        column_links.append({"source": "src_orders", "target": "iceberg_curated",
                             "column": cn, "masked": masked})

    return {
        "ok": True,
        "generated_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
        "active": {"generator": gen_active, "archive": arch_active},
        "stages": ["Sources", "CDC / Stream", "Compute", "Lake storage",
                   "Lakehouse", "Federation", "Serving"],
        "nodes": nodes,
        "edges": edges,
        "column_links": column_links,
        "om_connected": bool(OPENMETADATA_URL),
    }


def get_graph(use_cache: bool = True) -> dict[str, Any]:
    now = time.time()
    if use_cache and _cache["data"] and now - _cache["ts"] < _TTL:
        return _cache["data"]
    data = _build()
    _cache["data"] = data
    _cache["ts"] = now
    return data


@router.get("/datasets")
async def datasets() -> JSONResponse:
    return JSONResponse({"ok": True, "datasets": [
        {"key": d["key"], "label": d["label"], "engine": d["engine"],
         "color": d["color"], "table": d["fqtn"]}
        for d in DATASETS]})


@router.get("/graph")
async def graph(dataset: str | None = None, refresh: bool = False) -> JSONResponse:
    from starlette.concurrency import run_in_threadpool
    data = await run_in_threadpool(get_graph, not refresh)
    if dataset and dataset in DATASET_BY_KEY:
        # Filter to nodes reachable from / to the selected source, keep serving spine.
        keep = {f"src_{dataset}", f"topic_{dataset}", f"lake_{dataset}",
                "connect", "spark", "iceberg_curated", "iceberg_hadoop",
                "trino", "es", "rag", "vllm", "cc"}
        nodes = [n for n in data["nodes"] if n["id"] in keep]
        node_ids = {n["id"] for n in nodes}
        edges = [e for e in data["edges"] if e["source"] in node_ids and e["target"] in node_ids]
        out = {**data, "nodes": nodes, "edges": edges, "focus": dataset}
        return JSONResponse(out)
    return JSONResponse(data)