From 05794906a737ba4d298a2ece55ad3d6474612837 Mon Sep 17 00:00:00 2001
From: mo <mo@dell-atc.lan>
Date: Sat, 27 Jun 2026 11:44:16 +0200
Subject: [PATCH] perf(pii): native source-DB lookups for masking enforcement

Query postgres/mysql/mongo directly (fast, early LIMIT) instead of full Trino
table scans; Trino remains the path for the curated lakehouse + as a fallback.
---
 api/pii_catalog.py | 66 +++++++++++++++++++++++++++++++++++++---------
 1 file changed, 53 insertions(+), 13 deletions(-)

diff --git a/api/pii_catalog.py b/api/pii_catalog.py
index 9da5e2b..0711d80 100644
--- a/api/pii_catalog.py
+++ b/api/pii_catalog.py
@@ -40,13 +40,16 @@ _policy_cache: dict[str, bool] | None = None
 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"},
+     "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"},
+     "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"},
+     "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,
@@ -291,6 +294,44 @@ async def set_policy(body: MaskPolicyRequest) -> JSONResponse:
     return JSONResponse({"ok": True, "key": key, "columns": cols, "masked": bool(body.masked)})
 
 
+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.
@@ -307,19 +348,18 @@ async def lookup(body: LookupRequest) -> JSONResponse:
     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)
+        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: