feat(pii): self-service per-column masking policy enforced for the LLM
- 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.
This commit is contained in:
mo
+155
-6
@@ -8,14 +8,17 @@ 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"])
|
||||
|
||||
@@ -24,6 +27,14 @@ 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 = [
|
||||
@@ -42,6 +53,44 @@ DATASETS = [
|
||||
"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"),
|
||||
@@ -66,11 +115,6 @@ def _classify(col: str) -> str | None:
|
||||
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 = {
|
||||
@@ -150,9 +194,10 @@ def _build() -> dict[str, Any]:
|
||||
cat = heur or _OM_CAT.get(pii_tag, "PII")
|
||||
else:
|
||||
cat = heur
|
||||
masked = _is_masked(c, ds.get("masked_layer", False))
|
||||
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,
|
||||
})
|
||||
@@ -184,6 +229,110 @@ def get_pii(use_cache: bool = True) -> dict[str, Any]:
|
||||
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."),
|
||||
})
|
||||
|
||||
Reference in New Issue
Block a user