"""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 _policy_mtime: float = -1.0 # 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": "hadoop", "node_id": "iceberg_hadoop", "label": "Iceberg hadoop historical_sales_hdfs", "table": "iceberg.hadoop.historical_sales_hdfs", "table_name": "historical_sales_hdfs", "catalog": "iceberg", "schema": "hadoop", "om_fqn": "atc_trino.iceberg.hadoop.historical_sales_hdfs"}, {"key": "hdfs", "node_id": "hdfs", "label": "Hadoop HDFS historical_sales_hdfs", "table": "iceberg.hadoop.historical_sales_hdfs", "table_name": "historical_sales_hdfs", "catalog": "iceberg", "schema": "hadoop", "om_fqn": "atc_trino.iceberg.hadoop.historical_sales_hdfs"}, {"key": "cassandra", "node_id": "cassandra", "label": "Cassandra device_metrics", "table": "cassandra_telemetry.telemetry.device_metrics", "table_name": "device_metrics", "catalog": "cassandra_telemetry", "schema": "telemetry", "om_fqn": "atc_trino.cassandra_telemetry.telemetry.device_metrics", "native": {"engine": "cassandra"}}, {"key": "neo4j", "node_id": "neo4j", "label": "Neo4j Product/Supplier graph", "table": "neo4j_graph", "table_name": "graph", "catalog": "neo4j", "native": {"engine": "neo4j"}}, ] 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]: """Load policy from disk; re-read when the file mtime changes (Data Flow toggles).""" global _policy_cache, _policy_mtime try: mtime = POLICY_PATH.stat().st_mtime except Exception: mtime = 0.0 if _policy_cache is None or mtime != _policy_mtime: try: _policy_cache = {k: bool(v) for k, v in json.loads(POLICY_PATH.read_text()).items()} except Exception: if _policy_cache is None: _policy_cache = {} _policy_mtime = mtime return _policy_cache def _save_policy(p: dict[str, bool]) -> None: global _policy_cache, _policy_mtime _policy_cache = p try: POLICY_PATH.parent.mkdir(parents=True, exist_ok=True) POLICY_PATH.write_text(json.dumps(p, indent=2)) _policy_mtime = POLICY_PATH.stat().st_mtime except Exception: _policy_mtime = time.time() def invalidate_pii_caches() -> None: """Force catalog rebuild so chat/Data Flow see the same mask flags immediately.""" _cache["data"] = None _cache["ts"] = 0.0 try: import trino_federated as tf if hasattr(tf, "_dict_cache"): tf._dict_cache["data"] = None tf._dict_cache["at"] = 0.0 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$|customer_ip|client_ip", "IP"), (r"customer_id|client_id|user_id|account_id|member_id|device_id|subscriber_id|employee_id|guest_id|person_id|supplier_id", "IDENTIFIER"), (r"^payload$|payload|^notes$|^note$|free_text|freeform|raw_json|description", "FREEFORM"), ] _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 _neo4j_property_keys() -> list[str]: """Property keys across the Neo4j graph, used as the 'columns' for PII tagging.""" try: from neo4j import GraphDatabase uri = os.getenv("NEO4J_URI", f"bolt://{os.getenv('DB_HOST', '10.0.21.51')}:7687") drv = GraphDatabase.driver(uri, auth=(os.getenv("NEO4J_USER", "neo4j"), os.getenv("NEO4J_PASSWORD", "testpwd"))) with drv.session() as s: keys = [r["propertyKey"] for r in s.run("CALL db.propertyKeys() YIELD propertyKey RETURN propertyKey")] drv.close() return keys except Exception: return [] def _build() -> dict[str, Any]: datasets_out = [] total_pii = 0 total_masked = 0 om_used = False for ds in DATASETS: if (ds.get("native") or {}).get("engine") == "neo4j": cols = _neo4j_property_keys() else: 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) invalidate_pii_caches() # chat + Data Flow must reflect the toggle immediately return JSONResponse({"ok": True, "key": key, "columns": cols, "masked": bool(body.masked)}) def policy_column_lists(*, use_cache: bool = False) -> dict[str, Any]: """Live masked vs visible PII columns β€” same source as the Data Flow tab.""" catalog = get_pii(use_cache=use_cache) masked: list[dict[str, str]] = [] visible: list[dict[str, str]] = [] for d in catalog.get("datasets", []): key = d.get("key") or "" label = d.get("label") or key locked = bool(d.get("masked_layer") or d.get("policy_locked")) for c in d.get("pii_columns", []): entry = { "key": key, "label": label, "column": c.get("name") or "", "category": c.get("category") or "PII", "locked": locked, } (masked if c.get("masked") else visible).append(entry) summ = catalog.get("summary") or {} return { "catalog": catalog, "masked": masked, "visible": visible, "summary": summ, "mask_token": MASK_TOKEN, } def _interest_categories(message: str) -> list[str]: lower = (message or "").lower() cat_map = [ (("email", "e-mail", "mail"), "EMAIL"), (("phone", "telefoon", "mobile"), "PHONE"), (("name", "naam"), "NAME"), (("iban", "bank", "card"), "FINANCIAL"), (("ssn", "bsn", "national", "passport"), "NATIONAL_ID"), (("address", "adres"), "ADDRESS"), (("birth", "dob", "geboorte"), "DOB"), (("ip",), "IP"), ] out: list[str] = [] for words, cat in cat_map: if any(w in lower for w in words): out.append(cat) return out def _sample_rows_for_chat( *, categories: list[str] | None = None, max_datasets: int = 2, rows_per: int = 2, ) -> list[dict[str, Any]]: """Fetch a few live rows; values already policy-masked.""" catalog = get_pii(use_cache=False) prefer = ["mysql", "postgres", "mongodb", "cassandra", "neo4j", "curated"] samples: list[dict[str, Any]] = [] for key in prefer: dset = next((d for d in catalog.get("datasets", []) if d.get("key") == key), None) if not dset or not dset.get("pii_columns"): continue ds = DATASET_BY_KEY.get(key) if not ds: continue pii_cols = dset["pii_columns"] if categories: focus = [c for c in pii_cols if c.get("category") in categories] # Always keep one id-like visible column for context when focusing ids = [c for c in pii_cols if c.get("category") == "IDENTIFIER" and not c.get("masked")] pick = (ids[:1] + focus) if focus else pii_cols else: pick = pii_cols if not pick: continue # de-dupe preserving order seen: set[str] = set() select_cols: list[str] = [] masked_map: dict[str, bool] = {} for c in pick: name = c["name"] if name in seen: continue seen.add(name) select_cols.append(name) masked_map[name] = bool(c.get("masked")) if len(select_cols) >= 6: break name_col = next((c["name"] for c in pii_cols if c.get("category") == "NAME"), None) try: cols, rows = _lookup_rows(ds, select_cols, name_col, None, rows_per) except Exception: continue if not rows: continue rendered = [] for row in rows[:rows_per]: rendered.append({ cname: (MASK_TOKEN if masked_map.get(cname) else val) for cname, val in zip(cols, row) }) samples.append({ "key": key, "label": dset.get("label", key), "table": ds.get("table"), "rows": rendered, "masked_cols": [c for c, m in masked_map.items() if m], "visible_cols": [c for c, m in masked_map.items() if not m], }) if len(samples) >= max_datasets: break return samples def build_policy_evidence(*, max_datasets: int = 2, rows_per: int = 2) -> str: """Compact internal evidence for LLM context (not shown raw to users).""" snap = policy_column_lists(use_cache=False) masked = snap["masked"] visible = snap["visible"] summ = snap["summary"] lines = [ "=== PII POLICY (Data Flow synced) ===", f"Masked {summ.get('masked_columns', len(masked))}/{summ.get('pii_columns', 0)} Β· " f"visible {summ.get('unmasked_columns', len(visible))}. Token: {MASK_TOKEN}", "Masked: " + ", ".join(f"{m['key']}.{m['column']}" for m in masked[:25]) + ( f" …(+{len(masked)-25})" if len(masked) > 25 else "" ), "Visible: " + (", ".join(f"{v['key']}.{v['column']}" for v in visible[:25]) or "(none)"), ] for s in _sample_rows_for_chat(max_datasets=max_datasets, rows_per=rows_per): lines.append(f"Sample {s['label']}:") for row in s["rows"]: lines.append(" " + " | ".join(f"{k}={v}" for k, v in row.items())) return "\n".join(lines) def format_pii_chat_answer(message: str = "") -> str: """Short personal reply: masked β†’ say masked; visible β†’ show values. Never list field names.""" snap = policy_column_lists(use_cache=False) masked = snap["masked"] visible = snap["visible"] lower = (message or "").lower() interest = _interest_categories(message) greeting = any(w in lower for w in ("hi", "hello", "hey", "hallo", "goedemorgen", "goedemiddag")) hi = "Hi! " if greeting else "" label = { "EMAIL": "email", "PHONE": "phone number", "NAME": "name", "FINANCIAL": "bank / IBAN details", "NATIONAL_ID": "national ID", "ADDRESS": "address", "DOB": "date of birth", "IP": "IP address", } # No specific PII type asked β€” keep it vague, never enumerate columns if not interest: if any(w in lower for w in ("mask", "pii", "sensitive", "privacy", "personal")): return ( f"{hi}Personal data is protected by the masking policy. " f"Ask for something specific (an email, a phone number, a name…) and I'll tell you " f"whether I can share it β€” or only `{MASK_TOKEN}`." ) return ( f"{hi}I can't share personal data that's masked. " f"Ask me for an email, phone number, or name if you want to check." ) topic = ", ".join(label.get(c, c.lower()) for c in interest) interested_masked = [c for c in masked if c["category"] in interest] interested_visible = [c for c in visible if c["category"] in interest] # Collect visible sample *values* only (no column names in the reply) values: list[str] = [] if interested_visible: samples = _sample_rows_for_chat(categories=interest, max_datasets=2, rows_per=2) for s in samples: for row in s["rows"]: for col in interested_visible: if col["column"] in row: val = row[col["column"]] if val is None or val == "" or val == MASK_TOKEN: continue values.append(str(val)) # unique, preserve order seen: set[str] = set() uniq: list[str] = [] for v in values: if v not in seen: seen.add(v) uniq.append(v) values = uniq[:5] # Fully masked for this ask if interested_masked and not interested_visible: return ( f"{hi}No β€” that {topic} is masked (`{MASK_TOKEN}`). " "I can't share it." ) # Fully visible if interested_visible and not interested_masked: if values: listed = ", ".join(values) return f"{hi}Sure β€” here's what I can share: {listed}." return f"{hi}That {topic} isn't masked, but I don't have a sample value right now." # Mixed: some sources masked, some visible β€” still don't name columns if interested_visible and interested_masked: if values: listed = ", ".join(values) return ( f"{hi}Some of that is masked (`{MASK_TOKEN}`); " f"what I can share: {listed}." ) return ( f"{hi}Some of that {topic} is masked (`{MASK_TOKEN}`). " "I can't share the protected parts." ) return f"{hi}I don't have that personal data available." 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."), })