From 38e6461a8afbd7cb83709ddb998f03d40ba87755 Mon Sep 17 00:00:00 2001
From: Daniel Baumann <daniel.baumann@progress-linux.org>
Date: Tue, 23 Jan 2024 06:06:14 +0100
Subject: Merging upstream version 20.9.0.

Signed-off-by: Daniel Baumann <daniel.baumann@progress-linux.org>
---
 tests/test_optimizer.py | 30 ++++++++++++++++++++++++++++++
 1 file changed, 30 insertions(+)

(limited to 'tests/test_optimizer.py')

diff --git a/tests/test_optimizer.py b/tests/test_optimizer.py
index 9b68c78..3f96a0e 100644
--- a/tests/test_optimizer.py
+++ b/tests/test_optimizer.py
@@ -227,6 +227,15 @@ class TestOptimizer(unittest.TestCase):
 
     @patch("sqlglot.generator.logger")
     def test_qualify_columns(self, logger):
+        self.assertEqual(
+            optimizer.qualify_columns.qualify_columns(
+                parse_one("WITH x AS (SELECT a FROM db.y) SELECT * FROM db.x"),
+                schema={"db": {"x": {"z": "int"}, "y": {"a": "int"}}},
+                expand_stars=False,
+            ).sql(),
+            "WITH x AS (SELECT y.a AS a FROM db.y) SELECT * FROM db.x",
+        )
+
         self.assertEqual(
             optimizer.qualify_columns.qualify_columns(
                 parse_one("WITH x AS (SELECT a FROM db.y) SELECT z FROM db.x"),
@@ -285,6 +294,11 @@ class TestOptimizer(unittest.TestCase):
         schema = MappingSchema(self.schema, {"x": {"a"}, "y": {"b"}, "z": {"b"}})
         self.check_file("qualify_columns__with_invisible", qualify_columns, schema=schema)
 
+    def test_pushdown_cte_alias_columns(self):
+        self.check_file(
+            "pushdown_cte_alias_columns", optimizer.qualify_columns.pushdown_cte_alias_columns
+        )
+
     def test_qualify_columns__invalid(self):
         for sql in load_sql_fixtures("optimizer/qualify_columns__invalid.sql"):
             with self.subTest(sql):
@@ -534,6 +548,22 @@ FROM READ_CSV('tests/fixtures/optimizer/tpc-h/nation.csv.gz', 'delimiter', '|')
             level="warning",
         )
 
+    def test_struct_type_annotation(self):
+        tests = {
+            "SELECT STRUCT(1 AS col)": "STRUCT<col INT>",
+            "SELECT STRUCT(1 AS col, 2.5 AS row)": "STRUCT<col INT, row DOUBLE>",
+            "SELECT STRUCT(1)": "STRUCT<INT>",
+            "SELECT STRUCT(1 AS col, 2.5 AS row, struct(3.5 AS inner_col, 4 AS inner_row) AS nested_struct)": "STRUCT<col INT, row DOUBLE, nested_struct STRUCT<inner_col DOUBLE, inner_row INT>>",
+            "SELECT STRUCT(1 AS col, 2.5, ARRAY[1, 2, 3] AS nested_array, 'foo')": "STRUCT<col INT, DOUBLE, nested_array ARRAY<INT>, VARCHAR>",
+            "SELECT STRUCT(1, 2.5, 'bar')": "STRUCT<INT, DOUBLE, VARCHAR>",
+            'SELECT STRUCT(1 AS "CaseSensitive")': 'STRUCT<"CaseSensitive" INT>',
+        }
+
+        for sql, target_type in tests.items():
+            with self.subTest(sql):
+                expression = annotate_types(parse_one(sql))
+                assert expression.expressions[0].is_type(target_type)
+
     def test_literal_type_annotation(self):
         tests = {
             "SELECT 5": exp.DataType.Type.INT,
-- 
cgit v1.2.3