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
@@ -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"]}),
|
||||
|
||||
+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."),
|
||||
})
|
||||
|
||||
@@ -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<string | null>(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 ? (
|
||||
<div className="mt-1.5 border-t border-border pt-1.5">
|
||||
<div className="mb-1 flex items-center gap-1 text-[9px] font-medium text-rose-300">
|
||||
{selNode.pii.all_masked ? <ShieldCheck className="h-3 w-3 text-emerald-300" /> : <ShieldAlert className="h-3 w-3" />}
|
||||
{selNode.pii.pii_count} PII column{selNode.pii.pii_count === 1 ? '' : 's'}
|
||||
<div className="mb-1 flex items-center justify-between gap-1 text-[9px] font-medium text-rose-300">
|
||||
<span className="flex items-center gap-1">
|
||||
{selNode.pii.all_masked ? <ShieldCheck className="h-3 w-3 text-emerald-300" /> : <ShieldAlert className="h-3 w-3" />}
|
||||
{selNode.pii.pii_count} PII column{selNode.pii.pii_count === 1 ? '' : 's'}
|
||||
</span>
|
||||
{!selNode.pii.masked_layer && (
|
||||
<span className="flex gap-1">
|
||||
<button
|
||||
type="button"
|
||||
onClick={() => onToggleMask(selNode.pii!.key, '*', true)}
|
||||
className="rounded border border-emerald-400/40 px-1 text-[7px] text-emerald-300 hover:bg-emerald-500/15"
|
||||
>mask all</button>
|
||||
<button
|
||||
type="button"
|
||||
onClick={() => onToggleMask(selNode.pii!.key, '*', false)}
|
||||
className="rounded border border-rose-400/40 px-1 text-[7px] text-rose-300 hover:bg-rose-500/15"
|
||||
>unmask all</button>
|
||||
</span>
|
||||
)}
|
||||
</div>
|
||||
<p className="mb-1 text-[7px] leading-tight text-foreground-faint">
|
||||
Masked columns are hidden from the assistant — it refuses to reveal them.
|
||||
</p>
|
||||
<div className="flex flex-col gap-0.5">
|
||||
{selNode.pii.columns.map((c) => (
|
||||
<div key={c.name} className="flex items-center justify-between gap-1 text-[8px]">
|
||||
<span className="truncate font-mono text-foreground-muted">{c.name}</span>
|
||||
<span className="flex items-center gap-1">
|
||||
<span className="text-blue-300">{c.category}</span>
|
||||
<span className={c.masked ? 'text-emerald-300' : 'text-rose-300'}>
|
||||
{c.masked ? 'masked' : 'raw'}
|
||||
{selNode.pii.columns.map((c) => {
|
||||
const locked = c.policy_locked
|
||||
const busy = maskBusy === `${selNode.pii!.key}.${c.name}`
|
||||
return (
|
||||
<button
|
||||
key={c.name}
|
||||
type="button"
|
||||
disabled={locked || busy}
|
||||
onClick={() => onToggleMask(selNode.pii!.key, c.name, !c.masked)}
|
||||
title={locked ? 'Physically masked in the curated layer' : c.masked ? 'Masked — click to allow the assistant to read it' : 'Visible to the assistant — click to mask'}
|
||||
className={cn(
|
||||
'flex items-center justify-between gap-1 rounded px-1 py-0.5 text-[8px] transition-colors',
|
||||
locked ? 'cursor-default opacity-70' : 'hover:bg-surface-overlay',
|
||||
)}
|
||||
>
|
||||
<span className="truncate font-mono text-foreground-muted">{c.name}</span>
|
||||
<span className="flex items-center gap-1">
|
||||
<span className="text-blue-300">{c.category}</span>
|
||||
<span className={cn('inline-flex items-center gap-0.5', c.masked ? 'text-emerald-300' : 'text-rose-300')}>
|
||||
{busy ? <Loader2 className="h-2.5 w-2.5 animate-spin" />
|
||||
: c.masked ? <Lock className="h-2.5 w-2.5" /> : <LockOpen className="h-2.5 w-2.5" />}
|
||||
{c.masked ? 'masked' : 'visible'}
|
||||
</span>
|
||||
</span>
|
||||
</span>
|
||||
</div>
|
||||
))}
|
||||
</button>
|
||||
)
|
||||
})}
|
||||
</div>
|
||||
</div>
|
||||
) : (
|
||||
|
||||
@@ -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',
|
||||
|
||||
+2
-1
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user