Files
atc-agents/api/dataflow.py
T
mo 215ce111f1 feat(pii): self-service per-column masking policy enforced for the LLM
- pii_catalog: persistent per-column masking policy (default masked); GET/POST
  /api/pii/policy and POST /api/pii/lookup which redacts masked values server-side.
- get_pii masked flag now reflects the policy; dataflow exposes the dataset key.
- Data Flow PII inspector: per-column lock/unlock toggles + mask-all/unmask-all,
  so operators control exactly which data the assistant may reveal.
2026-06-27 11:36:36 +02:00

202 lines
8.6 KiB
Python

"""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 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.
NODES: list[dict[str, Any]] = [
{"id": "generator", "label": "Data Generator", "sub": "Airflow DAGs", "kind": "generator", "x": 10, "y": 50},
{"id": "hdfs", "label": "Hadoop HDFS", "sub": "historical_sales", "kind": "hadoop", "x": 10, "y": 88},
{"id": "postgres", "label": "PostgreSQL", "sub": "sales_orders", "kind": "source", "x": 30, "y": 18},
{"id": "mysql", "label": "MySQL", "sub": "employee_events", "kind": "source", "x": 30, "y": 44},
{"id": "mongodb", "label": "MongoDB", "sub": "events", "kind": "source", "x": 30, "y": 70},
{"id": "kafka", "label": "Kafka · Debezium", "sub": "CDC topics", "kind": "stream", "x": 52, "y": 40},
{"id": "s3_cdc", "label": "S3 CDC Archive", "sub": "object store", "kind": "sink", "x": 72, "y": 20},
{"id": "iceberg_hadoop", "label": "Iceberg · hadoop", "sub": "historical_sales_hdfs", "kind": "lakehouse", "x": 72, "y": 60},
{"id": "iceberg_curated", "label": "Iceberg · curated_masked", "sub": "masked PII", "kind": "lakehouse", "x": 72, "y": 86},
{"id": "trino", "label": "Trino", "sub": "query engine", "kind": "engine", "x": 90, "y": 52},
{"id": "openmetadata", "label": "OpenMetadata", "sub": "catalog · lineage · PII", "kind": "governance",
"x": 50, "y": 92, "url": "http://10.0.21.47:8585"},
]
# 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": "postgres", "to": "kafka", "kind": "cdc"},
{"from": "mysql", "to": "kafka", "kind": "cdc"},
{"from": "mongodb", "to": "kafka", "kind": "cdc"},
{"from": "kafka", "to": "s3_cdc", "kind": "archive"},
{"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": "trino", "to": "openmetadata", "kind": "catalog"},
]
_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
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
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 pii_catalog import get_pii
pii = get_pii()
except Exception:
pii = {"datasets": []}
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"
if n["id"] == "kafka":
metric = f"{cdc.get('window_total', 0)} chg/15m · {cdc.get('consumed', 0)} total"
node["level"] = "ok" if cdc.get("connected") else "warn"
elif n["id"] in ("postgres", "mysql", "mongodb"):
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"
# 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)
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"] == "cdc":
edge["active"] = cdc.get("by_source", {}).get(e["from"], 0) > 0
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")},
"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 = _build()
_cache["data"] = data
_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"})