Move tests to qualify_tables.sql

Author: georgesittas

sqlglot/optimizer/qualify_tables.py

@@ -131,7 +131,7 @@ def _set_alias(
                 _set_alias(source, target_alias=name or source.name or None, normalize=True)
 
                 source_fqn = ".".join(p.name for p in source.parts)
-                table_aliases[source_fqn] = exp.to_identifier(source.alias)
+                table_aliases[source_fqn] = source.args["alias"].this.copy()
 
                 if pivot:
                     target_alias = source.alias if pivot.unpivot else None
@@ -179,6 +179,6 @@ def _set_alias(
                 and (canonical_table := canonical_aliases.get(table, "")) != column.table
             ):
                 # Amend existing aliases, e.g. t.c -> _0.c if t is aliased to _0
-                column.set("table", exp.to_identifier(canonical_table))
+                column.set("table", exp.to_identifier(canonical_table, quoted=True))
 
     return expression

tests/fixtures/optimizer/qualify_tables.sql

@@ -213,5 +213,40 @@ WITH cte AS (SELECT 1 AS c, 'name' AS name) UPDATE t SET name = cte.name FROM ct
 WITH cte AS (SELECT 1 AS c, 'name' AS name) UPDATE c.db.t SET name = cte.name FROM cte WHERE cte.c = 1;
 
 # title: avoid qualifying CTE with DELETE
-WITH cte AS (SELECT 1 AS c, 'name' AS name) DELETE t FROM t AS t INNER JOIN cte ON t.id = cte.c
-WITH cte AS (SELECT 1 AS c, 'name' AS name) DELETE c.db.t FROM c.db.t AS t INNER JOIN cte ON t.id = cte.c
+WITH cte AS (SELECT 1 AS c, 'name' AS name) DELETE t FROM t AS t INNER JOIN cte ON t.id = cte.c;
+WITH cte AS (SELECT 1 AS c, 'name' AS name) DELETE c.db.t FROM c.db.t AS t INNER JOIN cte ON t.id = cte.c;
+
+# title: canonicalize single table alias
+# canonicalize_table_aliases: true
+SELECT * FROM t;
+SELECT * FROM c.db.t AS "_0";
+
+# title: canonicalize join table aliases
+# canonicalize_table_aliases: true
+SELECT * FROM t1 JOIN t2 ON t1.id = t2.id;
+SELECT * FROM c.db.t1 AS "_0" JOIN c.db.t2 AS "_1" ON "_0".id = "_1".id;
+
+# title: canonicalize join with different databases
+# canonicalize_table_aliases: true
+SELECT * FROM db1.users JOIN db2.users ON db1.users.id = db2.users.id;
+SELECT * FROM c.db1.users AS "_0" JOIN c.db2.users AS "_1" ON "_0".id = "_1".id;
+
+# title: canonicalize CTE alias
+# canonicalize_table_aliases: true
+WITH cte AS (SELECT * FROM t) SELECT * FROM cte;
+WITH cte AS (SELECT * FROM c.db.t AS "_0") SELECT * FROM cte AS "_1";
+
+# title: canonicalize subquery alias
+# canonicalize_table_aliases: true
+SELECT * FROM (SELECT * FROM t);
+SELECT * FROM (SELECT * FROM c.db.t AS "_0") AS "_1";
+
+# title: canonicalize multiple tables with subquery
+# canonicalize_table_aliases: true
+SELECT * FROM t1, (SELECT * FROM t2) AS sub, t3;
+SELECT * FROM c.db.t1 AS "_2", (SELECT * FROM c.db.t2 AS "_0") AS "_1", c.db.t3 AS "_3";
+
+# title: canonicalize CTE with PIVOT
+# canonicalize_table_aliases: true
+WITH cte AS (SELECT * FROM t) SELECT * FROM cte PIVOT(SUM(c) FOR v IN ('x', 'y'));
+WITH cte AS (SELECT * FROM c.db.t AS "_0") SELECT * FROM cte AS "_1" PIVOT(SUM(c) FOR v IN ('x', 'y')) AS "_2";

tests/test_optimizer.py

@@ -163,21 +163,25 @@ def check_file(
                 title = meta.get("title") or f"{i}, {sql}"
                 if only and title != only:
                     continue
+
                 dialect = meta.get("dialect")
                 leave_tables_isolated = meta.get("leave_tables_isolated")
                 validate_qualify_columns = meta.get("validate_qualify_columns")
+                canonicalize_table_aliases = meta.get("canonicalize_table_aliases")
+
+                func_kwargs = kwargs.copy()
 
-                func_kwargs = {**kwargs}
                 if leave_tables_isolated is not None:
                     func_kwargs["leave_tables_isolated"] = string_to_bool(leave_tables_isolated)
 
                 if validate_qualify_columns is not None:
                     func_kwargs["validate_qualify_columns"] = string_to_bool(
                         validate_qualify_columns
                     )
-
                 if dialect:
                     func_kwargs["dialect"] = dialect
+                if canonicalize_table_aliases:
+                    func_kwargs["canonicalize_table_aliases"] = canonicalize_table_aliases
 
                 future = pool.submit(parse_and_optimize, func, sql, dialect, **func_kwargs)
                 results[future] = (
@@ -240,77 +244,6 @@ def test_isolate_table_selects(self):
         )
 
     def test_qualify_tables(self):
-        self.assertEqual(
-            optimizer.qualify_tables.qualify_tables(
-                parse_one("SELECT * FROM t"),
-                db="db",
-                catalog="c",
-                canonicalize_table_aliases=True,
-            ).sql(),
-            'SELECT * FROM c.db.t AS "_0"',
-        )
-
-        self.assertEqual(
-            optimizer.qualify_tables.qualify_tables(
-                parse_one("SELECT * FROM t1 JOIN t2 ON t1.id = t2.id"),
-                db="db",
-                catalog="c",
-                canonicalize_table_aliases=True,
-            ).sql(),
-            'SELECT * FROM c.db.t1 AS "_0" JOIN c.db.t2 AS "_1" ON _0.id = _1.id',
-        )
-
-        self.assertEqual(
-            optimizer.qualify_tables.qualify_tables(
-                parse_one("SELECT * FROM db1.users JOIN db2.users ON db1.users.id = db2.users.id"),
-                catalog="c",
-                canonicalize_table_aliases=True,
-            ).sql(),
-            'SELECT * FROM c.db1.users AS "_0" JOIN c.db2.users AS "_1" ON _0.id = _1.id',
-        )
-
-        self.assertEqual(
-            optimizer.qualify_tables.qualify_tables(
-                parse_one("WITH cte AS (SELECT * FROM t) SELECT * FROM cte"),
-                db="db",
-                catalog="c",
-                canonicalize_table_aliases=True,
-            ).sql(),
-            'WITH cte AS (SELECT * FROM c.db.t AS "_0") SELECT * FROM cte AS "_1"',
-        )
-
-        self.assertEqual(
-            optimizer.qualify_tables.qualify_tables(
-                parse_one("SELECT * FROM (SELECT * FROM t)"),
-                db="db",
-                catalog="c",
-                canonicalize_table_aliases=True,
-            ).sql(),
-            'SELECT * FROM (SELECT * FROM c.db.t AS "_0") AS "_1"',
-        )
-
-        self.assertEqual(
-            optimizer.qualify_tables.qualify_tables(
-                parse_one("SELECT * FROM t1, (SELECT * FROM t2) AS sub, t3"),
-                db="db",
-                catalog="c",
-                canonicalize_table_aliases=True,
-            ).sql(),
-            'SELECT * FROM c.db.t1 AS "_2", (SELECT * FROM c.db.t2 AS "_0") AS "_1", c.db.t3 AS "_3"',
-        )
-
-        self.assertEqual(
-            optimizer.qualify_tables.qualify_tables(
-                parse_one(
-                    "WITH cte AS (SELECT * FROM t) SELECT * FROM cte PIVOT(SUM(c) FOR v IN ('x', 'y'))"
-                ),
-                db="db",
-                catalog="c",
-                canonicalize_table_aliases=True,
-            ).sql(),
-            'WITH cte AS (SELECT * FROM c.db.t AS "_0") SELECT * FROM cte AS "_1" PIVOT(SUM(c) FOR v IN (\'x\', \'y\')) AS "_2"',
-        )
-
         self.assertEqual(
             optimizer.qualify.qualify(
                 parse_one("WITH tesT AS (SELECT * FROM t1) SELECT * FROM test", "bigquery"),