Repository: tobymao/sqlglot

Issue:
[Snowflake][Column-level Lineage] Not able to parse lineage with PIVOT
**Fully reproducible code snippet**
SQL:
```sql
WITH source AS (
  SELECT 
    id
    , foo
    , bar
  FROM DB_NAME.SCHEMA_NAME.TABLE_NAME
),

enriched AS (
  SELECT * FROM source
  PIVOT(MAX(foo) 
          FOR bar 
          IN ('a', 'b', 'c', 'd', 'e'))
)

SELECT id, foo, bar FROM enriched
```

Python Code:
```python
with open("test_sql/test.sql", "r") as input:
    sql = input.read()

    parsed = sqlglot.parse_one(sql)
    parsed = sqlglot.optimizer.qualify.qualify(
        parsed,
        dialect="snowflake",
        catalog=None,
        db=None,
        qualify_columns=False,
        validate_qualify_columns=False,
        identify=False
    )

    sqlglot_db_schema = sqlglot.MappingSchema(dialect="snowflake", normalize=False)
    table = sqlglot.exp.Table(
        catalog=sqlglot.exp.Identifier(this="DB_NAME"),
        db=sqlglot.exp.Identifier(this="SCHEMA_NAME"),
        this=sqlglot.exp.Identifier(this="TABLE_NAME"),
    )
    table_schema = {
        "ID": "TEXT",
        "FOO": "NUMBER",
        "BAR": "TIMESTAMP_LTZ"
    }

    sqlglot_db_schema.add_table(
        table,
        column_mapping=table_schema,
    )
    
    lineage = sqlglot.lineage.lineage(
        "ID",
        parsed,
        dialect="snowflake",
        schema=sqlglot_db_schema
    )

    for node in lineage.walk():
        if node.downstream:
            pass
        elif isinstance(node.expression, sqlglot.exp.Table):
            table_ref = _TableName.from_sqlglot_table(node.expression)
            normalized_col = sqlglot.parse_one(node.name).this.name

            print(table_ref)
            print(normalized_col)
        else:
            pass
```

Expected output:
```
database=DB_NAME db_schema=SCHEMA_NAME table='TABLE_NAME'
ID
```

Actual output:
```
database=None db_schema=None table='SOURCE'
ID
```

**Official Documentation**
https://docs.snowflake.com/en/sql-reference/constructs/pivot



Discussion:
if you turn on validate_qualify_columns, you'll see that your code is incorrect.

```sql
WITH source AS (
  SELECT 1 id, 2 foo, 'a' bar
),

enriched AS (
  SELECT * FROM source
  PIVOT(MAX(foo) FOR bar  IN ('a', 'b', 'c', 'd', 'e'))
)

SELECT * FROM enriched
```

running this in snowflake yields columns id, a, b, c etc, the original foo, bar are no longer present. so your query is incorrect which is why lineage doesn't work
> running this in snowflake yields columns id, a, b, c etc, the original foo, bar are no longer present.

Adding to this, IIRC `foo` and `bar` become `"'foo'"` and `"'bar'"` respectively, i.e. Snowflake treats the names as strings and then converts them into identifiers by adding double-quotes. Perhaps you can get around this problem by changing your top-level query into:

```sql
SELECT id, "'foo'", "'bar'" FROM enriched
```
Hello @tobymao and @georgesittas , thank you so much for your reply. Apologies for confusing you but I would like to re-phrase my question a bit.

Please consider this runnable SQL query 
```sql
WITH source AS (
  SELECT 
    id
    , key
    , value
    , timestamp_1
    , timestamp_2
    FROM DB_NAME.SCHEMA_NAME.TABLE_NAME
),

enriched AS (
  SELECT * FROM source
  PIVOT(MAX(value) FOR key IN ('a', 'b', 'c'))
  AS final (id, timestamp_1, timestamp_2, col_1, col_2, col_3)
)

SELECT id, timestamp_1 FROM enriched
```

Ideally we should get `id` and `timestamp_1`'s lineage but from what I experimented it was not successful.

Based on my investigation, I found out that `sqlglot` will tend to analyse `source` in `enriched` CTE as a table. When I call `sqlglot.optimizer.qualify.qualify` with default db and schema as `CTE` and `CTE`, here is the result query:

```sql
WITH SOURCE AS (
  SELECT
    ID,
    KEY,
    VALUE,
    TIMESTAMP_1,
    TIMESTAMP_2
  FROM DB_NAME.SCHEMA_NAME.TABLE_NAME AS TABLE_NAME
), ENRICHED AS (
  SELECT
    *
  FROM CTE.CTE.SOURCE AS FINAL PIVOT(MAX(VALUE) FOR KEY IN ('a', 'b', 'c')) AS FINAL(ID, TIMESTAMP_1, TIMESTAMP_2, COL_1, COL_2, COL_3)
)
SELECT
  ID,
  TIMESTAMP_1
FROM ENRICHED
``` 
I think I see your point @billstark, thanks for clarifying. This may be a bug in how `PIVOT`s are qualified, seems like we're not using the table alias columns when expanding stars. I'll take a look tomorrow and let you know.
Hello @georgesittas , thanks for the quick response. Btw, nice avatar 😉 Samurai Champloo is the best!
Thanks - it is, indeed! :)
Hey @georgesittas , may I know that whether you got any ideas?
Hey - yep, I'm on this, just a bit busy but will fix it soon.

Resolving patch:
diff --git a/sqlglot/optimizer/pushdown_projections.py b/sqlglot/optimizer/pushdown_projections.py
--- a/sqlglot/optimizer/pushdown_projections.py
+++ b/sqlglot/optimizer/pushdown_projections.py
@@ -43,9 +43,8 @@ def pushdown_projections(expression, schema=None, remove_unused_selections=True)
         parent_selections = referenced_columns.get(scope, {SELECT_ALL})
         alias_count = source_column_alias_count.get(scope, 0)
 
-        if scope.expression.args.get("distinct") or (scope.parent and scope.parent.pivots):
-            # We can't remove columns SELECT DISTINCT nor UNION DISTINCT. The same holds if
-            # we select from a pivoted source in the parent scope.
+        # We can't remove columns SELECT DISTINCT nor UNION DISTINCT.
+        if scope.expression.args.get("distinct"):
             parent_selections = {SELECT_ALL}
 
         if isinstance(scope.expression, exp.Union):
@@ -78,7 +77,7 @@ def pushdown_projections(expression, schema=None, remove_unused_selections=True)
             # Push the selected columns down to the next scope
             for name, (node, source) in scope.selected_sources.items():
                 if isinstance(source, Scope):
-                    columns = selects.get(name) or set()
+                    columns = {SELECT_ALL} if scope.pivots else selects.get(name) or set()
                     referenced_columns[source].update(columns)
 
                 column_aliases = node.alias_column_names
diff --git a/sqlglot/optimizer/qualify_columns.py b/sqlglot/optimizer/qualify_columns.py
--- a/sqlglot/optimizer/qualify_columns.py
+++ b/sqlglot/optimizer/qualify_columns.py
@@ -386,7 +386,7 @@ def _expand_stars(
     pivot_exclude_columns = None
 
     pivot = t.cast(t.Optional[exp.Pivot], seq_get(scope.pivots, 0))
-    if isinstance(pivot, exp.Pivot):
+    if isinstance(pivot, exp.Pivot) and not pivot.alias_column_names:
         if pivot.unpivot:
             pivot_output_columns = [c.output_name for c in _unpivot_columns(pivot)]
 
@@ -431,12 +431,17 @@ def _expand_stars(
             table_id = id(table)
             columns_to_exclude = except_columns.get(table_id) or set()
 
-            if pivot and pivot_output_columns and pivot_exclude_columns:
-                implicit_columns = [c for c in columns if c not in pivot_exclude_columns]
-                if implicit_columns:
+            if pivot:
+                if pivot_output_columns and pivot_exclude_columns:
+                    pivot_columns = [c for c in columns if c not in pivot_exclude_columns]
+                    pivot_columns.extend(pivot_output_columns)
+                else:
+                    pivot_columns = pivot.alias_column_names
+
+                if pivot_columns:
                     new_selections.extend(
                         exp.alias_(exp.column(name, table=pivot.alias), name, copy=False)
-                        for name in implicit_columns + pivot_output_columns
+                        for name in pivot_columns
                         if name not in columns_to_exclude
                     )
                     continue
diff --git a/sqlglot/optimizer/qualify_tables.py b/sqlglot/optimizer/qualify_tables.py
--- a/sqlglot/optimizer/qualify_tables.py
+++ b/sqlglot/optimizer/qualify_tables.py
@@ -72,11 +72,15 @@ def qualify_tables(
                     if not source.args.get("catalog") and source.args.get("db"):
                         source.set("catalog", catalog)
 
+                pivots = pivots = source.args.get("pivots")
                 if not source.alias:
+                    # Don't add the pivot's alias to the pivoted table, use the table's name instead
+                    if pivots and pivots[0].alias == name:
+                        name = source.name
+
                     # Mutates the source by attaching an alias to it
                     alias(source, name or source.name or next_alias_name(), copy=False, table=True)
 
-                pivots = source.args.get("pivots")
                 if pivots and not pivots[0].alias:
                     pivots[0].set(
                         "alias", exp.TableAlias(this=exp.to_identifier(next_alias_name()))


Test patch:
diff --git a/tests/fixtures/optimizer/optimizer.sql b/tests/fixtures/optimizer/optimizer.sql
--- a/tests/fixtures/optimizer/optimizer.sql
+++ b/tests/fixtures/optimizer/optimizer.sql
@@ -643,6 +643,38 @@ FROM (
   FROM `u` AS `u` PIVOT(SUM(`u`.`f`) FOR `h` IN ('x', 'y'))
 ) AS `_q_0`;
 
+# title: selecting all columns from a pivoted source, pivot has column aliases
+# execute: false
+# dialect: snowflake
+WITH source AS (
+  SELECT
+    id,
+    key,
+    value,
+    timestamp_1,
+    timestamp_2
+  FROM DB_NAME.SCHEMA_NAME.TABLE_NAME
+),
+enriched AS (
+  SELECT * FROM source
+  PIVOT(MAX(value) FOR key IN ('a', 'b', 'c'))
+  AS final (id, timestamp_1, timestamp_2, col_1, col_2, col_3)
+)
+SELECT id, timestamp_1 FROM enriched;
+WITH "SOURCE" AS (
+  SELECT
+    "TABLE_NAME"."ID" AS "ID",
+    "TABLE_NAME"."KEY" AS "KEY",
+    "TABLE_NAME"."VALUE" AS "VALUE",
+    "TABLE_NAME"."TIMESTAMP_1" AS "TIMESTAMP_1",
+    "TABLE_NAME"."TIMESTAMP_2" AS "TIMESTAMP_2"
+  FROM "DB_NAME"."SCHEMA_NAME"."TABLE_NAME" AS "TABLE_NAME"
+)
+SELECT
+  "FINAL"."ID" AS "ID",
+  "FINAL"."TIMESTAMP_1" AS "TIMESTAMP_1"
+FROM "SOURCE" AS "SOURCE" PIVOT(MAX("SOURCE"."VALUE") FOR "SOURCE"."KEY" IN ('a', 'b', 'c')) AS "FINAL"("ID", "TIMESTAMP_1", "TIMESTAMP_2", "COL_1", "COL_2", "COL_3");
+
 # title: unpivoted table source with a single value column, unpivot columns can't be qualified
 # execute: false
 # dialect: snowflake
@@ -656,6 +688,23 @@ FROM "M_SALES" AS "M_SALES"("EMPID", "DEPT", "JAN", "FEB") UNPIVOT("SALES" FOR "
 ORDER BY
   "_q_0"."EMPID";
 
+# title: unpivoted table source, unpivot has column aliases
+# execute: false
+SELECT * FROM (SELECT * FROM m_sales) AS m_sales(empid, dept, jan, feb) UNPIVOT(sales FOR month IN (jan, feb)) AS unpiv(a, b, c, d);
+SELECT
+  "unpiv"."a" AS "a",
+  "unpiv"."b" AS "b",
+  "unpiv"."c" AS "c",
+  "unpiv"."d" AS "d"
+FROM (
+  SELECT
+    "m_sales"."empid" AS "empid",
+    "m_sales"."dept" AS "dept",
+    "m_sales"."jan" AS "jan",
+    "m_sales"."feb" AS "feb"
+  FROM "m_sales" AS "m_sales"
+) AS "m_sales" UNPIVOT("sales" FOR "month" IN ("m_sales"."jan", "m_sales"."feb")) AS "unpiv"("a", "b", "c", "d");
+
 # title: unpivoted derived table source with a single value column
 # execute: false
 # dialect: snowflake
diff --git a/tests/fixtures/optimizer/qualify_tables.sql b/tests/fixtures/optimizer/qualify_tables.sql
--- a/tests/fixtures/optimizer/qualify_tables.sql
+++ b/tests/fixtures/optimizer/qualify_tables.sql
@@ -31,6 +31,10 @@ SELECT (SELECT y.c FROM c.db.y AS y) FROM c.db.x AS x;
 SELECT * FROM x PIVOT (SUM(a) FOR b IN ('a', 'b'));
 SELECT * FROM c.db.x AS x PIVOT(SUM(a) FOR b IN ('a', 'b')) AS _q_0;
 
+# title: pivoted table, pivot has alias
+SELECT * FROM x PIVOT (SUM(a) FOR b IN ('a', 'b')) AS piv;
+SELECT * FROM c.db.x AS x PIVOT(SUM(a) FOR b IN ('a', 'b')) AS piv;
+
 # title: wrapped table without alias
 SELECT * FROM (tbl);
 SELECT * FROM (c.db.tbl AS tbl);