atc-agents/api/pii_catalog.py

"""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."),
    })