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.
# om_fqn = OpenMetadata table FQN (service.database.schema.table) for tag lookup.
DATASETS = [
    {"key": "postgres", "node_id": "postgres", "label": "PostgreSQL sales_orders",
     "table": "postgres_sales.public.sales_orders", "table_name": "sales_orders", "catalog": "postgres_sales",
     "om_fqn": "atc_postgres.postgres.public.sales_orders"},
    {"key": "mysql", "node_id": "mysql", "label": "MySQL employee_events",
     "table": "mysql_hr.hr.employee_events", "table_name": "employee_events", "catalog": "mysql_hr",
     "om_fqn": "atc_mysql.default.hr.employee_events"},
    {"key": "mongodb", "node_id": "mongodb", "label": "MongoDB events",
     "table": "mongodb_supplychain.supplychain.events", "table_name": "events", "catalog": "mongodb_supplychain",
     "om_fqn": "atc_mongodb.default.supplychain.events"},
    {"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,
     "om_fqn": "atc_trino.iceberg.curated_masked.sales_orders_masked"},
]

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


# OpenMetadata PII tag -> our category. OM applies PII.Sensitive / PII.NonSensitive
# plus optional General/PersonalData tags via auto-classification.
_OM_CAT = {
    "PII.Sensitive": "SENSITIVE",
    "PII.NonSensitive": "NON_SENSITIVE",
}


def _om_column_tags(fqn: str) -> dict[str, list[str]]:
    """Return {column_name: [tagFQN,...]} from OpenMetadata for a table FQN."""
    if not OPENMETADATA_URL:
        return {}
    url = f"{OPENMETADATA_URL}/api/v1/tables/name/{fqn}?fields=columns,tags"
    headers = {"Accept": "application/json"}
    if OPENMETADATA_TOKEN:
        headers["Authorization"] = f"Bearer {OPENMETADATA_TOKEN}"
    try:
        with httpx.Client(timeout=8.0) as client:
            r = client.get(url, headers=headers)
            if r.status_code != 200:
                return {}
            out: dict[str, list[str]] = {}
            for c in r.json().get("columns", []) or []:
                tags = [t.get("tagFQN") for t in (c.get("tags") or []) if t.get("tagFQN")]
                if tags:
                    out[c["name"]] = tags
            return out
    except Exception:
        return {}


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
    om_used = False
    for ds in DATASETS:
        cols = _trino_columns(ds["catalog"], ds.get("schema"), ds["table_name"])
        om_tags = _om_column_tags(ds["om_fqn"]) if ds.get("om_fqn") else {}
        if om_tags:
            om_used = True
        # Union of columns known via Trino and via OM (OM may exist before Trino sees it).
        all_cols = list(dict.fromkeys(cols + list(om_tags.keys())))
        pii_cols = []
        for c in all_cols:
            tags = om_tags.get(c, [])
            pii_tag = next((t for t in tags if t.startswith("PII.")), None)
            heur = _classify(c)
            if not pii_tag and not heur:
                continue
            # Prefer OM PII classification; enrich with heuristic category if present.
            if pii_tag:
                cat = heur or _OM_CAT.get(pii_tag, "PII")
            else:
                cat = heur
            masked = _is_masked(c, ds.get("masked_layer", False))
            pii_cols.append({
                "name": c, "category": cat, "masked": masked,
                "source": "openmetadata" if pii_tag else "heuristic",
                "om_tag": pii_tag,
            })
            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(all_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": "openmetadata+heuristic" if om_used else "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))