fix(optimizer): handle ColumnDef in typed table alias columns (#7542)

qualify_tables._set_alias called exp.to_identifier() on every column,
which rejects ColumnDef nodes produced by typed table-function signatures
like FROM fn() AS t(col INT). Route only strings through to_identifier
and copy any other expression (Identifier, ColumnDef) to preserve the
existing copy-before-reparent semantics.

Widens the columns/function_columns type annotations to include
exp.ColumnDef so mypyc's runtime type checks accept it.

Author: tobymao

sqlglot/optimizer/qualify_tables.py

@@ -80,7 +80,7 @@ def _set_alias(
         target_alias: str | None = None,
         scope: Scope | None = None,
         normalize: bool = False,
-        columns: Sequence[str | exp.Identifier] | None = None,
+        columns: Sequence[str | exp.Identifier | exp.ColumnDef] | None = None,
     ) -> None:
         alias = expression.args.get("alias") or exp.TableAlias()
 
@@ -97,7 +97,10 @@ def _set_alias(
         alias.set("this", exp.to_identifier(new_alias_name))
 
         if columns:
-            alias.set("columns", [exp.to_identifier(c) for c in columns])
+            alias.set(
+                "columns",
+                [exp.to_identifier(c) if isinstance(c, str) else c.copy() for c in columns],
+            )
 
         expression.set("alias", alias)
 
@@ -137,7 +140,7 @@ def _set_alias(
 
                 table_this = source.this
                 table_alias = source.args.get("alias")
-                function_columns: Sequence[str | exp.Identifier] | None = None
+                function_columns: Sequence[str | exp.Identifier | exp.ColumnDef] | None = None
                 if isinstance(table_this, exp.Func):
                     if not table_alias:
                         function_columns = ensure_list(

tests/fixtures/optimizer/qualify_tables.sql

@@ -281,3 +281,9 @@ SELECT g FROM GENERATE_SERIES(1, 2) AS t(g);
 # canonicalize_table_aliases: true
 SELECT g FROM GENERATE_SERIES(1,2) AS t(g);
 SELECT g FROM GENERATE_SERIES(1, 2) AS _0(g);
+
+# title: Qualify JSONB_TO_RECORDSET with typed alias columns and canonicalize_table_aliases
+# dialect: postgres
+# canonicalize_table_aliases: true
+SELECT jsonvalue.id FROM deal, JSONB_TO_RECORDSET(CAST(deal.type AS JSONB)) AS jsonvalue(id INT, key VARCHAR, text VARCHAR);
+SELECT _1.id FROM c.db.deal AS _0, JSONB_TO_RECORDSET(CAST(_0.type AS JSONB)) AS _1(id INT, key VARCHAR, text VARCHAR);

tests/test_optimizer.py

@@ -306,6 +306,19 @@ def test_qualify_tables(self):
             catalog="c",
         )
 
+    def test_qualify_tables_copies_typed_alias_columns(self):
+        expression = parse_one('SELECT * FROM JSON_TO_RECORDSET(z) AS y("rank" INT)')
+
+        original = expression.find(exp.Table).args["alias"].columns[0]
+        self.assertIsInstance(original, exp.ColumnDef)
+
+        optimizer.qualify_tables.qualify_tables(expression, canonicalize_table_aliases=True)
+
+        new = expression.find(exp.Table).args["alias"].columns[0]
+        self.assertIsInstance(new, exp.ColumnDef)
+        self.assertIsNot(original, new)
+        self.assertEqual(original.sql(), new.sql())
+
     def test_normalize(self):
         self.assertEqual(
             optimizer.normalize.normalize(