mo
215ce111f1
- 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.
339 lines
13 KiB
Python
339 lines
13 KiB
Python
"""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"},
|
|
{"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"},
|
|
]
|
|
|
|
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)})
|
|
|
|
|
|
@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))
|
|
col_sql = ", ".join(f'"{c}"' for c in select_cols)
|
|
|
|
where = ""
|
|
if body.search and name_col:
|
|
safe = body.search.replace("'", "''")
|
|
where = f" WHERE lower(cast(\"{name_col}\" AS varchar)) LIKE lower('%{safe}%')"
|
|
limit = max(1, min(body.limit or 5, 25))
|
|
sql = f"SELECT {col_sql} FROM {ds['table']}{where} LIMIT {limit}"
|
|
|
|
try:
|
|
cols, rows = _trino_query(sql)
|
|
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."),
|
|
})
|