ATC/atc-agents
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(pii): masking for Cassandra & Neo4j (graph nodes + live property read) + MongoDB payload/free-text

Browse Source
This commit is contained in:
mo
2026-06-27 21:10:16 +00:00
parent e2ac114863
commit cdc8aa4bf7
2 changed files with 40 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
api/pii_catalog.py
+28 -2
View File
@@ -62,6 +62,14 @@ DATASETS = [
"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}
@@ -112,7 +120,8 @@ PII_RULES: list[tuple[str, str]] = [
(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", "IDENTIFIER"),
(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}
@@ -182,13 +191,30 @@ def _trino_columns(catalog: str, schema: str | None, table_name: str) -> list[st
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:
cols = _trino_columns(ds["catalog"], ds.get("schema"), ds["table_name"])
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