atc-agents/api/pii_catalog.py

"""PII catalog for the Command Center Data Flow PII overlay.

Classifies columns of the known datasets into PII categories. Primary source is
OpenMetadata tags (when available, Phase 3); falls back to a name-based
heuristic over the live table columns (via Trino information_schema). Reports
per dataset which columns are PII and whether they are masked.
"""

from __future__ import annotations

import os
import re
import time
from typing import Any

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

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

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

# Datasets we surface in the PII overlay. node_id matches dataflow.py node ids.
DATASETS = [
    {"key": "postgres", "node_id": "postgres", "label": "PostgreSQL sales_orders",
     "table": "postgres_sales.public.sales_orders", "table_name": "sales_orders", "catalog": "postgres_sales"},
    {"key": "mysql", "node_id": "mysql", "label": "MySQL employee_events",
     "table": "mysql_hr.hr.employee_events", "table_name": "employee_events", "catalog": "mysql_hr"},
    {"key": "mongodb", "node_id": "mongodb", "label": "MongoDB events",
     "table": "mongodb_supplychain.supplychain.events", "table_name": "events", "catalog": "mongodb_supplychain"},
    {"key": "curated", "node_id": "iceberg_curated", "label": "Iceberg curated_masked",
     "table": "iceberg.curated_masked.sales_orders_masked", "table_name": "sales_orders_masked", "catalog": "iceberg",
     "schema": "curated_masked", "masked_layer": True},
]

# name fragment -> PII category
PII_RULES: list[tuple[str, str]] = [
    (r"email|e_mail", "EMAIL"),
    (r"phone|mobile|msisdn|tel", "PHONE"),
    (r"iban|account_no|bank|card|credit", "FINANCIAL"),
    (r"ssn|bsn|national_id|passport|tax_id", "NATIONAL_ID"),
    (r"first_name|last_name|full_name|customer_name|employee_name|contact_name|^name$", "NAME"),
    (r"address|street|city|zip|postal|postcode", "ADDRESS"),
    (r"dob|birth|date_of_birth", "DOB"),
    (r"ip_addr|ip_address", "IP"),
]

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


def _classify(col: str) -> str | None:
    c = col.lower()
    for pat, cat in PII_RULES:
        if re.search(pat, c):
            return cat
    return None


def _is_masked(col: str, dataset_masked_layer: bool) -> bool:
    c = col.lower()
    return dataset_masked_layer or c.endswith("_masked") or c.endswith("_hash") or c.endswith("_token")


def _trino_columns(catalog: str, schema: str | None, table_name: str) -> list[str]:
    sql = (
        f"SELECT column_name FROM {catalog}.information_schema.columns "
        f"WHERE table_name = '{table_name}'"
    )
    if schema:
        sql += f" AND table_schema = '{schema}'"
    try:
        with httpx.Client(timeout=10.0) as client:
            r = client.post(f"{TRINO_URL}/v1/statement", content=sql.encode(), headers={"X-Trino-User": TRINO_USER})
            d = r.json()
            rows: list[Any] = d.get("data") or []
            nxt = d.get("nextUri")
            while nxt:
                dd = client.get(nxt).json()
                rows += dd.get("data") or []
                if dd.get("error"):
                    break
                nxt = dd.get("nextUri")
            return [row[0] for row in rows]
    except Exception:
        return []


def _build() -> dict[str, Any]:
    datasets_out = []
    total_pii = 0
    total_masked = 0
    for ds in DATASETS:
        cols = _trino_columns(ds["catalog"], ds.get("schema"), ds["table_name"])
        pii_cols = []
        for c in cols:
            cat = _classify(c)
            if cat:
                masked = _is_masked(c, ds.get("masked_layer", False))
                pii_cols.append({"name": c, "category": cat, "masked": masked})
                total_pii += 1
                if masked:
                    total_masked += 1
        datasets_out.append({
            "key": ds["key"], "node_id": ds["node_id"], "label": ds["label"], "table": ds["table"],
            "exists": bool(cols), "masked_layer": ds.get("masked_layer", False),
            "pii_columns": pii_cols, "pii_count": len(pii_cols),
            "has_pii": bool(pii_cols),
            "all_masked": bool(pii_cols) and all(c["masked"] for c in pii_cols),
        })
    return {
        "ok": True, "source": "heuristic" if not OPENMETADATA_URL else "openmetadata+heuristic",
        "datasets": datasets_out,
        "summary": {"datasets": len(datasets_out), "pii_columns": total_pii, "masked_columns": total_masked,
                    "unmasked_columns": total_pii - total_masked},
    }


def get_pii(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("")
async def pii_overview(refresh: bool = False) -> JSONResponse:
    return JSONResponse(get_pii(use_cache=not refresh))