Merging upstream version 16.2.1.

Signed-off-by: Daniel Baumann <daniel.baumann@progress-linux.org>

1 files changed, 15 insertions, 2 deletions

diff --git a/tests/test_optimizer.py b/tests/test_optimizer.py
index 2ae6da9..94bd0ba 100644
--- a/tests/test_optimizer.py
+++ b/tests/test_optimizer.py
@@ -200,6 +200,15 @@ class TestOptimizer(unittest.TestCase):
     def test_qualify_columns(self):
         self.assertEqual(
             optimizer.qualify_columns.qualify_columns(
+                parse_one("WITH x AS (SELECT a FROM db.y) SELECT z FROM db.x"),
+                schema={"db": {"x": {"z": "int"}, "y": {"a": "int"}}},
+                infer_schema=False,
+            ).sql(),
+            "WITH x AS (SELECT y.a AS a FROM db.y) SELECT x.z AS z FROM db.x",
+        )
+
+        self.assertEqual(
+            optimizer.qualify_columns.qualify_columns(
                 parse_one("select y from x"),
                 schema={},
                 infer_schema=False,
@@ -544,9 +553,10 @@ FROM READ_CSV('tests/fixtures/optimizer/tpc-h/nation.csv.gz', 'delimiter', '|')
 
     def test_function_annotation(self):
         schema = {"x": {"cola": "VARCHAR", "colb": "CHAR"}}
-        sql = "SELECT x.cola || TRIM(x.colb) AS col FROM x AS x"
+        sql = "SELECT x.cola || TRIM(x.colb) AS col, DATE(x.colb) FROM x AS x"
 
-        concat_expr_alias = annotate_types(parse_one(sql), schema=schema).expressions[0]
+        expression = annotate_types(parse_one(sql), schema=schema)
+        concat_expr_alias = expression.expressions[0]
         self.assertEqual(concat_expr_alias.type.this, exp.DataType.Type.VARCHAR)
 
         concat_expr = concat_expr_alias.this
@@ -555,6 +565,9 @@ FROM READ_CSV('tests/fixtures/optimizer/tpc-h/nation.csv.gz', 'delimiter', '|')
         self.assertEqual(concat_expr.right.type.this, exp.DataType.Type.VARCHAR)  # TRIM(x.colb)
         self.assertEqual(concat_expr.right.this.type.this, exp.DataType.Type.CHAR)  # x.colb
 
+        date_expr = expression.expressions[1]
+        self.assertEqual(date_expr.type.this, exp.DataType.Type.DATE)
+
         sql = "SELECT CASE WHEN 1=1 THEN x.cola ELSE x.colb END AS col FROM x AS x"
 
         case_expr_alias = annotate_types(parse_one(sql), schema=schema).expressions[0]