diff --git a/api/dataflow.py b/api/dataflow.py index 0786377..988c584 100644 --- a/api/dataflow.py +++ b/api/dataflow.py @@ -143,6 +143,7 @@ def _build() -> dict[str, Any]: p = pii_by_node.get(n["id"]) if p: node["pii"] = { + "key": p["key"], "has_pii": p["has_pii"], "pii_count": p["pii_count"], "all_masked": p["all_masked"], "masked_layer": p.get("masked_layer", False), "categories": sorted({c["category"] for c in p["pii_columns"]}), diff --git a/api/pii_catalog.py b/api/pii_catalog.py index 904ba36..9da5e2b 100644 --- a/api/pii_catalog.py +++ b/api/pii_catalog.py @@ -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."), + }) diff --git a/ui/src/components/features/DataFlowView.tsx b/ui/src/components/features/DataFlowView.tsx index 9d16a83..8f7cd8a 100644 --- a/ui/src/components/features/DataFlowView.tsx +++ b/ui/src/components/features/DataFlowView.tsx @@ -1,7 +1,7 @@ import { useCallback, useEffect, useLayoutEffect, useMemo, useRef, useState } from 'react' -import { GitBranch, Play, RefreshCw, ShieldAlert, ShieldCheck, Loader2, Bot } from 'lucide-react' +import { GitBranch, Lock, LockOpen, Play, RefreshCw, ShieldAlert, ShieldCheck, Loader2, Bot } from 'lucide-react' import type { DataflowEdge, DataflowGraph, DataflowNode } from '../../types' -import { fetchDataflow, runDataflowMovement, toggleEtlAgent, fetchAgentOpsStatus } from '../../lib/api' +import { fetchDataflow, runDataflowMovement, toggleEtlAgent, fetchAgentOpsStatus, setPiiMask } from '../../lib/api' import { Badge } from '../ui/Badge' import { cn } from '../../lib/utils' @@ -158,6 +158,17 @@ export function DataFlowView() { setTimeout(() => setRefreshing(false), 400) }, [load]) + const [maskBusy, setMaskBusy] = useState(null) + const onToggleMask = useCallback(async (key: string, column: string, masked: boolean) => { + setMaskBusy(`${key}.${column}`) + try { + await setPiiMask(key, column, masked) + await load(true) + } finally { + setMaskBusy(null) + } + }, [load]) + const piiSummary = graph?.pii_summary ?? {} const movementEdges = useMemo( () => edges.filter((e) => e.movement_id), @@ -329,22 +340,57 @@ export function DataFlowView() { )} {selNode.pii?.has_pii ? (
-
- {selNode.pii.all_masked ? : } - {selNode.pii.pii_count} PII column{selNode.pii.pii_count === 1 ? '' : 's'} +
+ + {selNode.pii.all_masked ? : } + {selNode.pii.pii_count} PII column{selNode.pii.pii_count === 1 ? '' : 's'} + + {!selNode.pii.masked_layer && ( + + + + + )}
+

+ Masked columns are hidden from the assistant — it refuses to reveal them. +

- {selNode.pii.columns.map((c) => ( -
- {c.name} - - {c.category} - - {c.masked ? 'masked' : 'raw'} + {selNode.pii.columns.map((c) => { + const locked = c.policy_locked + const busy = maskBusy === `${selNode.pii!.key}.${c.name}` + return ( +
- ))} + + ) + })}
) : ( diff --git a/ui/src/lib/api.ts b/ui/src/lib/api.ts index 8a66385..6244d11 100644 --- a/ui/src/lib/api.ts +++ b/ui/src/lib/api.ts @@ -155,6 +155,14 @@ export async function fetchPii(refresh = false): Promise<{ datasets: PiiDataset[ ) } +export function setPiiMask(key: string, column: string, masked: boolean) { + return fetch('/api/pii/policy', { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ key, column, masked }), + }) +} + export function toggleEtlAgent(enabled?: boolean) { return fetch('/api/agent-ops/etl/toggle', { method: 'POST', diff --git a/ui/src/types.ts b/ui/src/types.ts index 800f99f..552665d 100644 --- a/ui/src/types.ts +++ b/ui/src/types.ts @@ -70,7 +70,7 @@ export type CdcStats = { consumed: number } -export type PiiColumn = { name: string; category: string; masked: boolean } +export type PiiColumn = { name: string; category: string; masked: boolean; policy_locked?: boolean } export type DataflowNode = { id: string @@ -83,6 +83,7 @@ export type DataflowNode = { metric: string | null url?: string pii?: { + key: string has_pii: boolean pii_count: number all_masked: boolean