Fix: Transpile bug with pivot subquery string literal column names

Author: fivetran-amrutabhimsenayachit

sqlglot-integration-tests

@@ -2492,7 +2492,7 @@ def _pivot_in_value_aliases(self, expression: exp.Pivot) -> list[exp.Expression]
 
         # Derive the per-value suffix from the first stored column vs the first IN-list value.
         # This correctly handles dialects (e.g. Spark single-agg) that ignore agg aliases.
-        first_base = in_exprs[0].alias_or_name
+        first_base = in_exprs[0].sql() if src_identify_pivot_strings else in_exprs[0].alias_or_name
         first_stored = columns[0].name
 
         # exit if only suffix matches, not prefix. (e.g. BigQuery, which cannot be fixed)

sqlglot/generator.py

@@ -5270,7 +5270,12 @@ def _parse_pivot(self) -> exp.Pivot | None:
                 if isinstance(seq_get(pivot_field_expressions, 0), exp.PivotAny):
                     continue
 
-                all_fields.append([fld.alias_or_name for fld in pivot_field_expressions])
+                all_fields.append(
+                    [
+                        fld.sql() if self.IDENTIFY_PIVOT_STRINGS else fld.alias_or_name
+                        for fld in pivot_field_expressions
+                    ]
+                )
 
             if all_fields:
                 if names:

sqlglot/parser.py

@@ -830,7 +830,7 @@ def test_duckdb(self):
             },
         )
         self.validate_all(
-            "SELECT * FROM produce PIVOT(SUM(sales) FOR quarter IN ('Q1', 'Q2'))",
+            "SELECT * FROM produce PIVOT(SUM(sales) FOR quarter IN ('Q1' AS \"'Q1'\", 'Q2' AS \"'Q2'\"))",
             read={
                 "snowflake": "SELECT * FROM produce PIVOT(SUM(produce.sales) FOR produce.quarter IN ('Q1', 'Q2'))",
             },

tests/dialects/test_duckdb.py

@@ -712,19 +712,19 @@ def test_spark(self):
             },
         )
         self.validate_all(
-            "SELECT piv.Q1 FROM (SELECT * FROM produce PIVOT(SUM(sales) FOR quarter IN ('Q1', 'Q2'))) AS piv",
+            "SELECT piv.Q1 FROM (SELECT * FROM produce PIVOT(SUM(sales) FOR quarter IN ('Q1' AS `'Q1'`, 'Q2' AS `'Q2'`))) AS piv",
             read={
                 "snowflake": "SELECT piv.Q1 FROM produce PIVOT(SUM(sales) FOR quarter IN ('Q1', 'Q2')) piv",
             },
         )
         self.validate_all(
-            "SELECT piv.Q1 FROM (SELECT * FROM (SELECT * FROM produce) PIVOT(SUM(sales) FOR quarter IN ('Q1', 'Q2'))) AS piv",
+            "SELECT piv.Q1 FROM (SELECT * FROM (SELECT * FROM produce) PIVOT(SUM(sales) FOR quarter IN ('Q1' AS `'Q1'`, 'Q2' AS `'Q2'`))) AS piv",
             read={
                 "snowflake": "SELECT piv.Q1 FROM (SELECT * FROM produce) PIVOT(SUM(sales) FOR quarter IN ('Q1', 'Q2')) piv",
             },
         )
         self.validate_all(
-            "SELECT * FROM produce PIVOT(SUM(produce.sales) FOR quarter IN ('Q1', 'Q2'))",
+            "SELECT * FROM produce PIVOT(SUM(produce.sales) FOR quarter IN ('Q1' AS `'Q1'`, 'Q2' AS `'Q2'`))",
             read={
                 "snowflake": "SELECT * FROM produce PIVOT (SUM(produce.sales) FOR produce.quarter IN ('Q1', 'Q2'))",
             },

tests/dialects/test_spark.py

@@ -639,8 +639,8 @@ PIVOT(SUM(`u_cte`.`f`) AS `sum` FOR `u_cte`.`h` IN ('x', 'y')) AS `_0`;
 SELECT * FROM u PIVOT (SUM(f) FOR h IN ('x', 'y'));
 SELECT
   "_0"."G" AS "G",
-  "_0"."X" AS "X",
-  "_0"."Y" AS "Y"
+  "_0"."'x'" AS "'x'",
+  "_0"."'y'" AS "'y'"
 FROM "U" AS "U"
 PIVOT(SUM("U"."F") FOR "U"."H" IN ('x', 'y')) AS "_0";
 

tests/fixtures/optimizer/optimizer.sql

@@ -688,7 +688,7 @@ def test_pivot_columns(self):
                 "bigquery": ["prop", "rudder"],
                 "duckdb": ["prop", "rudder"],
                 "redshift": ["prop", "rudder"],
-                "snowflake": ["prop", "rudder"],
+                "snowflake": ['''"'prop'"''', '''"'rudder'"'''],
                 "spark": ["prop", "rudder"],
             },
             everything_aliased: {