DRC-3809: resolve PIVOT output columns to real source columns (CLL)

recce/util/cll.py

@@ -265,6 +265,17 @@ def reset(self):
         self.other_error_nodes = 0
 
 
+def _normalize_pivot_name(name: str) -> str:
+    """Normalize a PIVOT value/column name for matching.
+
+    sqlglot names generated pivot columns differently across dialects and
+    versions (e.g. ``JAN`` for an explicit projection vs. ``'jan'`` for a
+    ``select *`` expansion), so strip surrounding quotes and upper-case before
+    comparing. Matching is structural — never against the synthetic alias.
+    """
+    return name.strip("'\"").upper()
+
+
 def _dedeup_depends_on(depends_on: List[CllColumnDep]) -> List[CllColumnDep]:
     # deduplicate the depends_on list
     dedup_set = set()
@@ -327,9 +338,10 @@ def _cll_select_scope(scope: Scope, scope_cll_map: dict[Scope, CllResult]) -> Cl
     table_alias_map = {t.alias_or_name: t.name for t in scope.tables}
     select = scope.expression
 
-    def source_column_dependency(ref_column: exp.Column) -> Optional[CllColumn]:
-        column_name = ref_column.name
-        table_name = ref_column.table if ref_column.table != "" else next(iter(table_alias_map.values()))
+    def _resolve_source_column(column_name: str, table_name: str) -> Optional[CllColumn]:
+        """Resolve a column against a real entry in ``scope.sources`` (a base
+        table or a sub-Scope). Returns None when ``table_name`` is not a
+        registered source — e.g. a PIVOT's synthetic alias, handled below."""
         source = scope.sources.get(table_name, None)  # transformation_type: exp.Table | Scope
         if isinstance(source, Scope):
             ref_cll_result = _resolve_scope_cll(source, scope_cll_map)
@@ -346,6 +358,72 @@ def source_column_dependency(ref_column: exp.Column) -> Optional[CllColumn]:
         else:
             return None
 
+    # Per-scope PIVOT map (DRC-3809). `qualify` points a pivoted query's
+    # projections at the pivot's table alias (`P`, or a synthesized `_0` /
+    # `_q_0` when unaliased), but sqlglot never registers that alias in
+    # `scope.sources` — the `Pivot` hangs off the base table's
+    # `.args["pivots"]`. Without help, every projected pivot column falls
+    # through to the phantom-alias fallback below, losing both column- and
+    # model-level lineage. Build a map keyed by the pivot alias so those
+    # projections resolve back to the real pre-pivot source columns. Keyed off
+    # the pivot node's *structure*, never the alias string (version-dependent).
+    pivot_map: Dict[str, dict] = {}
+    for base_source in scope.sources.values():
+        if not isinstance(base_source, exp.Table):
+            continue
+        for pivot in base_source.args.get("pivots") or []:
+            alias = pivot.alias
+            if not alias:
+                continue
+            # Source columns the pivot consumes: the measure aggregate(s) and
+            # the pivoted-for dimension column(s). Both already reference the
+            # base alias, so they resolve through the normal source logic.
+            consumed_deps: List[CllColumnDep] = []
+            for measure in pivot.expressions:
+                for col in measure.find_all(exp.Column):
+                    resolved = _resolve_source_column(col.name, col.table)
+                    if resolved is not None:
+                        consumed_deps.extend(resolved.depends_on)
+            value_names: set = set()
+            for field in pivot.args.get("fields") or []:
+                if not isinstance(field, exp.In):
+                    continue
+                for col in field.this.find_all(exp.Column):
+                    resolved = _resolve_source_column(col.name, col.table)
+                    if resolved is not None:
+                        consumed_deps.extend(resolved.depends_on)
+                for value in field.expressions:
+                    value_names.add(_normalize_pivot_name(value.output_name))
+            pivot_map[alias] = {
+                "base": base_source,
+                "consumed_deps": _dedeup_depends_on(consumed_deps),
+                "value_names": value_names,
+            }
+
+    def source_column_dependency(ref_column: exp.Column) -> Optional[CllColumn]:
+        column_name = ref_column.name
+        table_name = ref_column.table if ref_column.table != "" else next(iter(table_alias_map.values()))
+        resolved = _resolve_source_column(column_name, table_name)
+        if resolved is not None:
+            return resolved
+        pivot = pivot_map.get(table_name)
+        if pivot is not None:
+            if _normalize_pivot_name(column_name) in pivot["value_names"]:
+                # Generated value column: derived from the measure + dimension
+                # source columns the pivot consumes.
+                return CllColumn(
+                    name=column_name,
+                    transformation_type="derived",
+                    depends_on=list(pivot["consumed_deps"]),
+                )
+            # Pass-through column: copied unchanged from the base relation.
+            return CllColumn(
+                name=column_name,
+                transformation_type="passthrough",
+                depends_on=[CllColumnDep(node=pivot["base"].name, column=column_name)],
+            )
+        return None
+
     def subquery_cll(subquery: exp.Subquery) -> Optional[CllResult]:
         select = subquery.find(exp.Select)
         if select is None:
@@ -494,6 +572,12 @@ def selected_column_dependency(ref_column: exp.Column) -> Optional[CllColumn]:
                 elif selected_column_dependency(ref_column) is not None:
                     m2c.extend(selected_column_dependency(ref_column).depends_on)
 
+    # PIVOT (DRC-3809): the measure + dimension columns a pivot consumes are
+    # model-level source provenance, analogous to group-by / where columns.
+    # Without this, a pivoted query loses ALL model-level lineage.
+    for pivot in pivot_map.values():
+        m2c.extend(pivot["consumed_deps"])
+
     for source in scope.sources.values():
         if not isinstance(source, Scope):
             continue

tests/util/test_cll.py

@@ -885,3 +885,53 @@ def test_correlated_subquery_in_join_condition(self):
         assert_column(result, "customer_name", "passthrough", [("customers", "customer_name")])
         assert_column(result, "first_order_id", "renamed", [("orders", "order_id")])
         assert_column(result, "first_order_date", "renamed", [("orders", "ordered_at")])
+
+    def test_pivot(self):
+        """PIVOT (DRC-3809): generated value columns must resolve to their
+        pre-pivot source columns, not a phantom pivot-alias node.
+
+        sqlglot points every projected column at the pivot's table alias
+        (``P``, or a synthesized ``_0`` / ``_q_0`` when unaliased), which it
+        never registers in ``scope.sources``. The fix keys off the pivot
+        node's structure, so it is immune to that alias-string drift across
+        sqlglot versions.
+        """
+        schema = {"t1": {"id": "int", "month": "varchar", "rev": "int"}}
+
+        # AC-1 + AC-2: Snowflake explicit-alias. qualify upper-cases snowflake
+        # identifiers, so expected node/column names are upper-case.
+        sql = """
+        select id, jan, feb from t1 pivot (sum(rev) for month in ('jan','feb')) as p
+        """
+        result = cll(sql, schema=schema, dialect="snowflake")
+        # AC-1: model-level lineage preserved to the columns the PIVOT consumes.
+        assert_model(result, [("T1", "REV"), ("T1", "MONTH")])
+        # AC-2: value columns are derived from measure+dimension; ID passes through.
+        assert_column(result, "JAN", "derived", [("T1", "REV"), ("T1", "MONTH")])
+        assert_column(result, "FEB", "derived", [("T1", "REV"), ("T1", "MONTH")])
+        assert_column(result, "ID", "passthrough", [("T1", "ID")])
+        # No phantom pivot-alias node anywhere.
+        m2c, c2c = result
+        assert all(d.node not in ("P", "_0", "_Q_0", "_q_0") for d in m2c), m2c
+        for col in c2c.values():
+            assert all(d.node not in ("P", "_0", "_Q_0", "_q_0") for d in col.depends_on), col
+
+        # select * (no explicit alias): guards the synthesized-alias path.
+        # Same expected lineage; the output value column keys are the pivot
+        # literals, but they must still resolve to T1.REV / T1.MONTH.
+        sql = """
+        select * from t1 pivot (sum(rev) for month in ('jan','feb'))
+        """
+        result = cll(sql, schema=schema, dialect="snowflake")
+        assert_model(result, [("T1", "REV"), ("T1", "MONTH")])
+        assert_column(result, "ID", "passthrough", [("T1", "ID")])
+        m2c, c2c = result
+        assert all(d.node not in ("P", "_0", "_Q_0", "_q_0") for d in m2c), m2c
+        # The two generated value columns (keyed by the pivot literals) are
+        # derived from the measure + dimension source columns, phantom-free.
+        value_cols = [c for name, c in c2c.items() if name.upper() != "ID"]
+        assert len(value_cols) == 2, list(c2c.keys())
+        for col in value_cols:
+            assert col.transformation_type == "derived", col
+            deps = {(d.node, d.column) for d in col.depends_on}
+            assert deps == {("T1", "REV"), ("T1", "MONTH")}, col