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 json
import os
import re
import time
from pathlib import Path
from typing import Any

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

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

# User-controlled masking policy. Each PII column can be masked (hidden from the
# LLM) or unmasked. Persisted to disk so it survives restarts. The default is
# "masked" — sensitive data is withheld until an operator explicitly opts out.
POLICY_PATH = Path(os.getenv("MASKING_POLICY_PATH", "/data/masking_policy.json"))
DEFAULT_MASKED = True
MASK_TOKEN = "🔒 MASKED (masking policy ON)"
_policy_cache: dict[str, bool] | None = None

# 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",
     "native": {"engine": "postgres", "table": "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",
     "native": {"engine": "mysql", "table": "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",
     "native": {"engine": "mongo", "db": "supplychain", "coll": "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"},
]

KEY_BY_NODE = {ds["node_id"]: ds["key"] for ds in DATASETS}
DATASET_BY_KEY = {ds["key"]: ds for ds in DATASETS}


def _load_policy() -> dict[str, bool]:
    global _policy_cache
    if _policy_cache is None:
        try:
            _policy_cache = {k: bool(v) for k, v in json.loads(POLICY_PATH.read_text()).items()}
        except Exception:
            _policy_cache = {}
    return _policy_cache


def _save_policy(p: dict[str, bool]) -> None:
    global _policy_cache
    _policy_cache = p
    try:
        POLICY_PATH.parent.mkdir(parents=True, exist_ok=True)
        POLICY_PATH.write_text(json.dumps(p, indent=2))
    except Exception:
        pass


def _resolve_key(key_or_node: str) -> str | None:
    if key_or_node in DATASET_BY_KEY:
        return key_or_node
    return KEY_BY_NODE.get(key_or_node)


def is_masked(key: str, column: str, *, masked_layer: bool = False) -> bool:
    """Resolve whether a column is masked: physical masked layer is always masked;
    otherwise the user policy override wins, falling back to DEFAULT_MASKED."""
    if masked_layer:
        return True
    return _load_policy().get(f"{key}.{column}", DEFAULT_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


# 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(ds["key"], c, masked_layer=ds.get("masked_layer", False))
            pii_cols.append({
                "name": c, "category": cat, "masked": masked,
                "policy_locked": ds.get("masked_layer", False),
                "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


def _trino_query(sql: str, timeout: float = 20.0) -> tuple[list[str], list[list[Any]]]:
    cols: list[str] = []
    rows: list[list[Any]] = []
    with httpx.Client(timeout=timeout) as client:
        d = client.post(f"{TRINO_URL}/v1/statement", content=sql.encode(), headers={"X-Trino-User": TRINO_USER}).json()
        while True:
            if d.get("error"):
                raise RuntimeError(d["error"].get("message", "trino error"))
            c = d.get("columns")
            if c and not cols:
                cols = [x["name"] for x in c]
            rows += d.get("data") or []
            nxt = d.get("nextUri")
            if not nxt:
                break
            d = client.get(nxt).json()
    return cols, rows


class MaskPolicyRequest(BaseModel):
    key: str            # dataset key or node id
    column: str         # column name, or "*" for every PII column of the dataset
    masked: bool


class LookupRequest(BaseModel):
    key: str            # dataset key or node id
    search: str | None = None
    limit: int = 5


@router.get("")
async def pii_overview(refresh: bool = False) -> JSONResponse:
    return JSONResponse(get_pii(use_cache=not refresh))


@router.get("/policy")
async def get_policy() -> JSONResponse:
    return JSONResponse({"ok": True, "default_masked": DEFAULT_MASKED, "policy": _load_policy()})


@router.post("/policy")
async def set_policy(body: MaskPolicyRequest) -> JSONResponse:
    key = _resolve_key(body.key)
    if not key:
        return JSONResponse({"ok": False, "error": f"unknown dataset {body.key}"}, status_code=400)
    data = get_pii()
    dset = next((d for d in data["datasets"] if d["key"] == key), None)
    if body.column == "*":
        cols = [c["name"] for c in (dset["pii_columns"] if dset else [])]
    else:
        cols = [body.column]
    if not cols:
        return JSONResponse({"ok": False, "error": "no columns to update"}, status_code=400)
    p = dict(_load_policy())
    for col in cols:
        p[f"{key}.{col}"] = bool(body.masked)
    _save_policy(p)
    _cache["data"] = None  # force rebuild so masked flags reflect the new policy
    return JSONResponse({"ok": True, "key": key, "columns": cols, "masked": bool(body.masked)})


def _lookup_rows(ds: dict[str, Any], select: list[str], name_col: str | None,
                 search: str | None, limit: int) -> tuple[list[str], list[list[Any]]]:
    """Fetch rows from the source. Native DB queries (fast, early LIMIT) for
    postgres/mysql/mongo; Trino for the curated lakehouse table."""
    from sql_console import _run_postgres, _run_mysql, _mongo_client  # local import avoids cycle

    nat = ds.get("native") or {}
    engine = nat.get("engine")
    safe = (search or "").replace("'", "''")

    if engine == "postgres":
        cols_sql = ", ".join(f'"{c}"' for c in select)
        where = f' WHERE "{name_col}" ILIKE \'%{safe}%\'' if (search and name_col) else ""
        res = _run_postgres(f"SELECT {cols_sql} FROM {nat['table']}{where} LIMIT {limit}", limit=limit)
        return res["columns"], res["rows"]
    if engine == "mysql":
        cols_sql = ", ".join(f"`{c}`" for c in select)
        where = f" WHERE `{name_col}` LIKE '%{safe}%'" if (search and name_col) else ""
        res = _run_mysql(f"SELECT {cols_sql} FROM {nat['table']}{where} LIMIT {limit}", limit=limit)
        return res["columns"], res["rows"]
    if engine == "mongo":
        cli = _mongo_client()
        try:
            coll = cli[nat.get("db", "supplychain")][nat["coll"]]
            filt = {name_col: {"$regex": safe, "$options": "i"}} if (search and name_col) else {}
            proj = {c: 1 for c in select}
            proj["_id"] = 0
            docs = list(coll.find(filt, proj).limit(limit))
        finally:
            cli.close()
        return select, [[d.get(c) for c in select] for d in docs]

    # Trino (curated lakehouse) or fallback
    col_sql = ", ".join(f'"{c}"' for c in select)
    where = f" WHERE lower(cast(\"{name_col}\" AS varchar)) LIKE lower('%{safe}%')" if (search and name_col) else ""
    return _trino_query(f"SELECT {col_sql} FROM {ds['table']}{where} LIMIT {limit}")


@router.post("/lookup")
async def lookup(body: LookupRequest) -> JSONResponse:
    """Look up actual records in a dataset, enforcing the masking policy server-side.
    Masked column values are replaced with MASK_TOKEN and never leave this process."""
    key = _resolve_key(body.key)
    if not key:
        return JSONResponse({"ok": False, "error": f"unknown dataset {body.key}"}, status_code=400)
    ds = DATASET_BY_KEY[key]
    dset = next((d for d in get_pii()["datasets"] if d["key"] == key), None)
    if not dset or not dset["pii_columns"]:
        return JSONResponse({"ok": False, "error": "no PII columns known for this dataset"}, status_code=400)

    pii_cols = dset["pii_columns"]
    masked_map = {c["name"]: c["masked"] for c in pii_cols}
    name_col = next((c["name"] for c in pii_cols if c["category"] == "NAME"), None)
    select_cols = list(dict.fromkeys(c["name"] for c in pii_cols))
    limit = max(1, min(body.limit or 5, 25))

    try:
        cols, rows = _lookup_rows(ds, select_cols, name_col, body.search, limit)
    except Exception:  # noqa: BLE001
        try:  # native failed → fall back to Trino over the same table
            col_sql = ", ".join(f'"{c}"' for c in select_cols)
            safe = (body.search or "").replace("'", "''")
            where = f" WHERE lower(cast(\"{name_col}\" AS varchar)) LIKE lower('%{safe}%')" if (body.search and name_col) else ""
            cols, rows = _trino_query(f"SELECT {col_sql} FROM {ds['table']}{where} LIMIT {limit}")
        except Exception as exc:  # noqa: BLE001
            return JSONResponse({"ok": False, "error": f"query failed: {exc}"}, status_code=502)

    out_rows = []
    for row in rows:
        rec: dict[str, Any] = {}
        for cname, val in zip(cols, row):
            rec[cname] = MASK_TOKEN if masked_map.get(cname) else val
        out_rows.append(rec)

    return JSONResponse({
        "ok": True, "key": key, "table": ds["table"],
        "rows": out_rows, "row_count": len(out_rows),
        "masked_columns": [c for c, m in masked_map.items() if m],
        "unmasked_columns": [c for c, m in masked_map.items() if not m],
        "note": ("Values shown as the mask token are withheld by the active masking policy "
                 "and were never sent to the model."),
    })