PII: obfuscate patterns

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "wlens"
-version = "0.1.1"
+version = "0.2.0"
 description = "Warehouse lens for AI agents — per-table markdown docs + read-only SQL adaptor, zero-runtime schema exploration for Claude Code, Cursor, Continue."
 readme = "README.md"
 license = { file = "LICENSE" }

src/wlens/adapters/dbt.py

@@ -21,6 +21,7 @@
 
 import json
 import logging
+import re
 from dataclasses import dataclass
 from pathlib import Path
 
@@ -173,13 +174,24 @@ def _build_test_index(manifest: dict) -> dict[tuple[str, str | None], list[str]]
            - "accepted_values(values=['shipped','completed','pending'])"
     """
     idx: dict[tuple[str, str | None], list[str]] = {}
+    sources_by_name = {
+        f"{node['source_name']}.{node['name']}": node_id
+        for node_id, node in manifest.get("sources", {}).items()
+        if "source_name" in node and "name" in node
+    }
+    models_by_name = {
+        node["name"]: node_id
+        for node_id, node in manifest.get("nodes", {}).items()
+        if node.get("resource_type") == "model" and "name" in node
+    }
+
     for node in manifest.get("nodes", {}).values():
         if node.get("resource_type") != "test":
             continue
-        attached = node.get("attached_node")
+        attached = node.get("attached_node") or _resolve_attached_from_metadata(
+            node, sources_by_name, models_by_name
+        )
         if not attached:
-            # Older manifests: derive from refs/sources. The first model/source
-            # in depends_on.nodes is usually the target.
             deps = (node.get("depends_on") or {}).get("nodes") or []
             attached = next(
                 (d for d in deps if d.startswith(("model.", "source."))),
@@ -188,12 +200,37 @@ def _build_test_index(manifest: dict) -> dict[tuple[str, str | None], list[str]]
         if not attached:
             continue
         idx.setdefault((attached, node.get("column_name")), []).append(_format_test(node))
-    # Stable ordering makes markdown diffs boring.
     for tests in idx.values():
         tests.sort()
     return idx
 
 
+_MODEL_TEMPLATE_SOURCE = re.compile(r"source\(\s*'([^']+)'\s*,\s*'([^']+)'\s*\)")
+_MODEL_TEMPLATE_REF = re.compile(r"ref\(\s*'([^']+)'\s*\)")
+
+
+def _resolve_attached_from_metadata(
+    test_node: dict,
+    sources_by_name: dict[str, str],
+    models_by_name: dict[str, str],
+) -> str | None:
+    """For tests with `attached_node: None`, parse `test_metadata.kwargs.model`
+    to find the node where the test is defined. The kwarg looks like
+    `{{ get_where_subquery(source('acme', 'Loss_Payment')) }}` or
+    `{{ get_where_subquery(ref('my_model')) }}`.
+    """
+    template = (test_node.get("test_metadata") or {}).get("kwargs", {}).get("model")
+    if not isinstance(template, str):
+        return None
+    match = _MODEL_TEMPLATE_SOURCE.search(template)
+    if match:
+        return sources_by_name.get(f"{match.group(1)}.{match.group(2)}")
+    match = _MODEL_TEMPLATE_REF.search(template)
+    if match:
+        return models_by_name.get(match.group(1))
+    return None
+
+
 def _format_test(test_node: dict) -> str:
     """Render one test node as a compact grep-friendly string."""
     meta = test_node.get("test_metadata") or {}

src/wlens/config.py

@@ -77,6 +77,9 @@ class OutputConfig:
     dir: str = DEFAULT_OUTPUT_DIR
     include_sample_rows: bool = True
     sample_size: int = DEFAULT_SAMPLE_SIZE
+    # Extra value-level obfuscation rules appended to the built-in defaults
+    # (email, UUID, URL, IP, phone). Each item: {pattern, replacement}.
+    obfuscate: list[dict] = field(default_factory=list)
 
 
 @dataclass
@@ -159,6 +162,7 @@ def _build_config(raw: dict, repo_root: Path) -> Config:
         dir=output_raw.get("dir", DEFAULT_OUTPUT_DIR),
         include_sample_rows=bool(output_raw.get("include_sample_rows", True)),
         sample_size=int(output_raw.get("sample_size", DEFAULT_SAMPLE_SIZE)),
+        obfuscate=list(output_raw.get("obfuscate") or []),
     )
 
     entities: list[EntityConfig] = []

src/wlens/render/markdown.py

@@ -13,12 +13,14 @@
 
 from __future__ import annotations
 
+import functools
 import logging
 from dataclasses import dataclass
 from pathlib import Path
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Callable
 
 from ..adapters.base import Column, Entity, Parent
+from .obfuscate import compile_rules, obfuscate
 from .pii import PII_REDACTED, pii_column_names
 from .preserve import read_manual_block
 
@@ -48,10 +50,15 @@ def render_and_write_all(
 
     inline_map = _inline_map(custom_entities)
     sample_size = config.output.sample_size if config.output.include_sample_rows else 0
+    obfuscate_value = functools.partial(
+        obfuscate, rules=compile_rules(config.output.obfuscate)
+    )
 
     for entity in entities:
         head_rows = _fetch_head(executor, entity, sample_size) if executor else []
-        body = _render_entity(entity, head_rows, inline_map.get(entity.slug))
+        body = _render_entity(
+            entity, head_rows, inline_map.get(entity.slug), obfuscate_value
+        )
         _write_file(output_dir / entity.filename, body)
 
     index_path = output_dir / "_index.md"
@@ -67,6 +74,7 @@ def _render_entity(
     entity: Entity,
     head_rows: list[dict],
     custom_entity: "TableCatalog | None",
+    obfuscate_value: Callable[[str], str],
 ) -> str:
     kind_label = "model" if entity.kind == "model" else "source"
     if custom_entity is not None:
@@ -80,7 +88,9 @@ def _render_entity(
     if custom_entity is not None:
         lines.extend(custom_entity.render())
     lines.extend(_render_parents(entity.parents))
-    lines.extend(_render_sample_rows(head_rows, pii_column_names(entity.columns)))
+    lines.extend(
+        _render_sample_rows(head_rows, pii_column_names(entity.columns), obfuscate_value)
+    )
     lines.extend(_render_compiled_sql(entity.compiled_sql))
     return "\n".join(lines).rstrip() + "\n"
 
@@ -115,7 +125,11 @@ def _render_parents(parents: list[Parent]) -> list[str]:
     return lines
 
 
-def _render_sample_rows(head_rows: list[dict], pii_columns: set[str]) -> list[str]:
+def _render_sample_rows(
+    head_rows: list[dict],
+    pii_columns: set[str],
+    obfuscate_value: Callable[[str], str],
+) -> list[str]:
     """Column-first (transposed) rendering: one bullet per column, pipe-separated values."""
     if not head_rows:
         return []
@@ -126,7 +140,7 @@ def _render_sample_rows(head_rows: list[dict], pii_columns: set[str]) -> list[st
         if c in pii_columns:
             values = " | ".join([PII_REDACTED] * n_rows)
         else:
-            values = " | ".join(_cell(row.get(c)) for row in head_rows)
+            values = " | ".join(_cell(row.get(c), obfuscate_value) for row in head_rows)
         lines.append(f"- `{c}`: {values}")
     lines.append("")
     return lines
@@ -236,10 +250,12 @@ def _truncate(value) -> str | None:
     return s
 
 
-def _cell(v) -> str:
+def _cell(v, obfuscate_value: Callable[[str], str]) -> str:
     if v is None:
         return "NULL"
     s = str(v).replace("|", "\\|").replace("\n", " ")
+    # Obfuscate before truncation so long emails/URLs aren't cut mid-pattern.
+    s = obfuscate_value(s)
     if len(s) > CELL_MAX_LEN:
         s = s[:CELL_MAX_LEN] + "…"
     return s

src/wlens/render/obfuscate.py

@@ -0,0 +1,85 @@
+"""Value-level PII obfuscation for rendered sample-row cells.
+
+Complement to render/pii.py (which redacts whole columns by name). This
+module runs after column-level redaction and scrubs PII shapes embedded
+inside otherwise-innocuous values — e.g. an email pasted into a workspace
+name, or a UUID inside a description field.
+
+Defaults cover only the shapes that are always problematic: email,
+phone, IP. URLs and UUIDs are deliberately excluded — UUIDs are warehouse
+primary keys everywhere, URLs are not always PII, and scrubbing either
+costs more in readability than it gains. Projects can append extra
+patterns (including UUID or URL rules) via `output.obfuscate` in
+wlens.yml.
+"""
+
+from __future__ import annotations
+
+import re
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True)
+class ObfuscationRule:
+    name: str
+    pattern: re.Pattern[str]
+    replacement: str
+
+
+# IP before phone: `1.2.3.4` should render as `<ip>`, not be partially matched.
+DEFAULT_RULES: tuple[ObfuscationRule, ...] = (
+    ObfuscationRule(
+        name="email",
+        pattern=re.compile(r"[\w.+-]+@[\w-]+(?:\.[\w-]+)+"),
+        replacement="<email>",
+    ),
+    ObfuscationRule(
+        name="ip",
+        pattern=re.compile(r"\b(?:\d{1,3}\.){3}\d{1,3}\b"),
+        replacement="<ip>",
+    ),
+    # Phone is split into two rules so neither matches DD-MM-YYYY / MM-DD-YYYY
+    # date strings. Real phones have ≥10 digits; common date formats have ≤8.
+    ObfuscationRule(
+        # International with explicit +country-code: `+1 555-123-4567`,
+        # `+44 20 7946 0958`. The leading `+` is the disambiguator.
+        name="phone_intl",
+        pattern=re.compile(
+            r"\+\d{1,3}[\s.-]?\d{1,4}[\s.-]?\d{2,9}[\s.-]?\d{2,9}"
+        ),
+        replacement="<phone>",
+    ),
+    ObfuscationRule(
+        # North American 3-3-4 shape: `(555) 123-4567`, `555-123-4567`,
+        # `555.123.4567`. The strict 3-3-4 grouping rules out date formats,
+        # whose leading day/month group is 1-2 digits.
+        name="phone_na",
+        pattern=re.compile(r"\(?\d{3}\)?[\s.-]\d{3}[\s.-]\d{4}"),
+        replacement="<phone>",
+    ),
+)
+
+
+def compile_rules(extras: list[dict] | None) -> tuple[ObfuscationRule, ...]:
+    """Return DEFAULT_RULES plus any user-supplied extras (appended)."""
+    rules: list[ObfuscationRule] = list(DEFAULT_RULES)
+    for i, extra in enumerate(extras or []):
+        if not isinstance(extra, dict) or "pattern" not in extra or "replacement" not in extra:
+            raise ValueError(
+                f"output.obfuscate[{i}] must have `pattern` and `replacement`. Got: {extra!r}"
+            )
+        rules.append(
+            ObfuscationRule(
+                name=str(extra.get("name") or f"extra_{i}"),
+                pattern=re.compile(str(extra["pattern"])),
+                replacement=str(extra["replacement"]),
+            )
+        )
+    return tuple(rules)
+
+
+def obfuscate(value: str, rules: tuple[ObfuscationRule, ...] = DEFAULT_RULES) -> str:
+    """Apply rules in order; return the scrubbed string."""
+    for rule in rules:
+        value = rule.pattern.sub(rule.replacement, value)
+    return value

src/wlens/templates/wlens.yml

@@ -15,3 +15,18 @@ executor:
   #   database: ${WLENS_DB_NAME}
   #   user:     ${WLENS_DB_USER}
   #   password: ${WLENS_DB_PASSWORD}
+
+# Value-level scrubbing of sample-row cells.
+# Defaults always on: email, IPv4, phone → <email> / <ip> / <phone>.
+# Append project-specific patterns here. Common additions:
+# output:
+#   obfuscate:
+#     # UUIDs (warehouse primary keys — opt in if you don't want them in docs)
+#     - pattern: '\b[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}\b'
+#       replacement: '<uuid>'
+#     # Stripe object IDs (cus_, sub_, in_, pi_, ch_, pm_, prod_, price_, inv_)
+#     - pattern: '\b(cus|sub|in|pi|ch|pm|prod|price|inv)_[A-Za-z0-9]+\b'
+#       replacement: '<stripe_id>'
+#     # Project-specific identifier example
+#     - pattern: '\bemp-\d{6}\b'
+#       replacement: '<employee_id>'

tests/test_render_markdown.py

@@ -29,6 +29,66 @@ def _setup_project(tmp_path: Path, dbt_project: Path, extra: str = "") -> Path:
     return cfg_path
 
 
+class _StubExecutor:
+    """Minimal `run_direct(sql)` stand-in for sample-row fetching."""
+
+    def __init__(self, headers: list[str], rows: list[tuple]):
+        self._headers = headers
+        self._rows = rows
+
+    def run_direct(self, sql: str):  # noqa: ARG002 — sql ignored, fixed payload
+        return self._headers, self._rows, False
+
+
+def test_sample_rows_obfuscate_embedded_pii(tmp_path, dbt_project):
+    uuid_pattern = (
+        r"\b[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-"
+        r"[0-9a-fA-F]{4}-[0-9a-fA-F]{12}\b"
+    )
+    cfg_path = tmp_path / "wlens.yml"
+    cfg_path.write_text(
+        textwrap.dedent(f"""
+            adapter:
+              kind: dbt
+              project_dir: {dbt_project.relative_to(tmp_path)}
+              include_prefixes: [dim_]
+            output:
+              dir: .claude/schema
+              include_sample_rows: true
+              sample_size: 2
+              obfuscate:
+                - pattern: '{uuid_pattern}'
+                  replacement: '<uuid>'
+                - pattern: '\\bemp-\\d{{6}}\\b'
+                  replacement: '<employee_id>'
+        """).lstrip()
+    )
+    cfg = load_config(cfg_path)
+    entities = [e for e in DbtAdapter(cfg).list_entities() if e.slug == "prod.dim_widget"]
+    assert entities, "expected dim_widget in fixtures"
+
+    executor = _StubExecutor(
+        headers=["widget_id", "color", "notes"],
+        rows=[
+            (1, "red", "owner foo@bar.com id 550e8400-e29b-41d4-a716-446655440000"),
+            (2, "blue", "ticket from 10.0.0.5 by emp-123456"),
+        ],
+    )
+    render_and_write_all(entities, [], executor, cfg)
+
+    md = (cfg.output_dir / "prod.dim_widget.md").read_text()
+    # Built-in patterns scrubbed.
+    assert "foo@bar.com" not in md
+    assert "10.0.0.5" not in md
+    assert "<email>" in md
+    assert "<ip>" in md
+    # User-supplied extras applied.
+    assert "550e8400" not in md
+    assert "<uuid>" in md
+    assert "emp-123456" not in md
+    assert "<employee_id>" in md
+
+
 def test_renders_one_file_per_entity(tmp_path, dbt_project):
     cfg = load_config(_setup_project(tmp_path, dbt_project))
     entities = DbtAdapter(cfg).list_entities()