fix(optimizer)!: qualify UNPIVOT on CTE sources (#7550)

Two fixes that together let `SELECT * FROM cte UNPIVOT(...)` qualify
  correctly in BigQuery and other dialects.

  - parser: `UNPIVOT`'s pre-`FOR` value column(s) and the `FOR` field are now
    parsed as `Identifier` (was `Column`). Those names don't reference
    existing columns — they're new output names. The `IN`-list items stay as
    `Column` since they do reference source-table columns. `PIVOT` is
    unchanged.

  - optimizer/resolver: when a CTE is pivoted, the scope stores an
    `exp.Table` under the pivot alias rather than the CTE's `Scope`, so
    column resolution couldn't see the CTE's columns. Fall back to
    `scope.cte_sources` in that case. Guarded on `not source.db` and
    `source.args.get("pivots")` so a real `db.x` that happens to share a
    name with a CTE doesn't misroute through the CTE scope.

  - optimizer/qualify_columns: removes the post-hoc filter in
    `validate_qualify_columns` that excluded unpivot output names from
    `scope.unqualified_columns` — they're `Identifier`s now and never land
    there. `_unpivot_columns` yields `Identifier`s; `output_name` still
    works.

  Test: new assertion in `test_qualify_columns` covering the CTE + `UNPIVOT` shape.

Author: georgesittas

sqlglot/optimizer/qualify_columns.py

@@ -138,13 +138,6 @@ def validate_qualify_columns(expression: E, sql: str | None = None) -> E:
 
                 raise OptimizeError(error_msg)
 
-            if unqualified_columns and scope.pivots and scope.pivots[0].unpivot:
-                # New columns produced by the UNPIVOT can't be qualified, but there may be columns
-                # under the UNPIVOT's IN clause that can and should be qualified. We recompute
-                # this list here to ensure those in the former category will be excluded.
-                unpivot_columns = set(_unpivot_columns(scope.pivots[0]))
-                unqualified_columns = [c for c in unqualified_columns if c not in unpivot_columns]
-
             all_unqualified_columns.extend(unqualified_columns)
 
     if all_unqualified_columns:
@@ -188,13 +181,18 @@ def _separate_pseudocolumns(scope: Scope, pseudocolumns: set[str]) -> None:
         scope.clear_cache()
 
 
-def _unpivot_columns(unpivot: exp.Pivot) -> Iterator[exp.Column]:
+def _unpivot_columns(unpivot: exp.Pivot) -> Iterator[exp.Identifier]:
     name_columns = [
         field.this
         for field in unpivot.fields
-        if isinstance(field, exp.In) and isinstance(field.this, exp.Column)
+        if isinstance(field, exp.In) and isinstance(field.this, exp.Identifier)
     ]
-    value_columns = (c for e in unpivot.expressions for c in e.find_all(exp.Column))
+    value_columns = (
+        ident
+        for e in unpivot.expressions
+        for ident in (e.expressions if isinstance(e, exp.Tuple) else [e])
+        if isinstance(ident, exp.Identifier)
+    )
 
     return itertools.chain(name_columns, value_columns)
 

sqlglot/optimizer/resolver.py

@@ -139,6 +139,17 @@ def get_source_columns(self, name: str, only_visible: bool = False) -> Sequence[
 
             source = self.scope.sources[name]
 
+            # A pivoted CTE reference is stored as an exp.Table in the scope sources (see
+            # _traverse_tables in scope.py), but the underlying CTE Scope still holds the
+            # column information we need to resolve pre-pivot columns.
+            if (
+                isinstance(source, exp.Table)
+                and not source.db
+                and source.args.get("pivots")
+                and source.name in self.scope.cte_sources
+            ):
+                source = self.scope.cte_sources[source.name]
+
             if isinstance(source, exp.Table):
                 columns = self.schema.column_names(source, only_visible)
             elif isinstance(source, Scope) and isinstance(

sqlglot/parser.py

@@ -275,6 +275,15 @@ def _resolve_dialect(dialect: DialectType) -> Dialect:
     return Dialect.get_or_raise(dialect)
 
 
+def _unpivot_target(expr: exp.Expr) -> exp.Expr:
+    # UNPIVOT's pre-FOR values and FOR field are new output names, not column references.
+    if isinstance(expr, exp.Column) and not expr.table:
+        return expr.this
+    if isinstance(expr, exp.Tuple):
+        expr.set("expressions", [_unpivot_target(e) for e in expr.expressions])
+    return expr
+
+
 SENTINEL_NONE: Token = Token(TokenType.SENTINEL, "SENTINEL")
 
 
@@ -5142,6 +5151,12 @@ def _parse_pivot(self) -> exp.Pivot | None:
             )
         )
 
+        if unpivot:
+            pivot.set("expressions", [_unpivot_target(e) for e in pivot.expressions])
+            for pivot_field in pivot.fields:
+                if isinstance(pivot_field, exp.In):
+                    pivot_field.set("this", _unpivot_target(pivot_field.this))
+
         if not self._match_set((TokenType.PIVOT, TokenType.UNPIVOT), advance=False):
             pivot.set("alias", self._parse_table_alias())
 

tests/test_optimizer.py

@@ -632,6 +632,21 @@ def test_qualify_columns(self, logger):
             "SELECT t.end AS end FROM t AS t",
         )
 
+        self.assertEqual(
+            optimizer.qualify.qualify(
+                parse_one(
+                    "WITH produce AS (SELECT 'Kale' AS product, 51 AS q1, 23 AS q2) "
+                    "SELECT * FROM produce UNPIVOT(sales FOR quarter IN (q1, q2))",
+                    dialect="bigquery",
+                ),
+                dialect="bigquery",
+            ).sql(dialect="bigquery"),
+            "WITH `produce` AS (SELECT 'Kale' AS `product`, 51 AS `q1`, 23 AS `q2`) "
+            "SELECT `produce`.`product` AS `product`, `produce`.`quarter` AS `quarter`, "
+            "`produce`.`sales` AS `sales` FROM `produce` AS `produce` "
+            "UNPIVOT(`sales` FOR `quarter` IN (`produce`.`q1`, `produce`.`q2`)) AS `produce`",
+        )
+
     def test_validate_columns(self):
         with self.assertRaisesRegex(
             OptimizeError, "Column 'foo' could not be resolved. Line: 1, Col: 10"