Adding upstream version 22.2.0.upstream/22.2.0

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

1 files changed, 47 insertions, 1 deletions

diff --git a/tests/test_optimizer.py b/tests/test_optimizer.py
index af8c3cd..046e5a6 100644
--- a/tests/test_optimizer.py
+++ b/tests/test_optimizer.py
@@ -205,6 +205,7 @@ class TestOptimizer(unittest.TestCase):
             optimizer.qualify_tables.qualify_tables,
             db="db",
             catalog="c",
+            set_dialect=True,
         )
 
     def test_normalize(self):
@@ -285,6 +286,15 @@ class TestOptimizer(unittest.TestCase):
             "SELECT `test`.`bar_bazfoo_$id` AS `bar_bazfoo_$id` FROM `test` AS `test`",
         )
 
+        qualified = optimizer.qualify.qualify(
+            parse_one("WITH t AS (SELECT 1 AS c) (SELECT c FROM t)")
+        )
+        self.assertIs(qualified.selects[0].parent, qualified.this)
+        self.assertEqual(
+            qualified.sql(),
+            'WITH "t" AS (SELECT 1 AS "c") (SELECT "t"."c" AS "c" FROM "t" AS "t")',
+        )
+
         self.check_file(
             "qualify_columns", qualify_columns, execute=True, schema=self.schema, set_dialect=True
         )
@@ -348,6 +358,23 @@ class TestOptimizer(unittest.TestCase):
         self.assertEqual("CONCAT('a', x, 'bc')", simplified_concat.sql(dialect="presto"))
         self.assertEqual("CONCAT('a', x, 'bc')", simplified_safe_concat.sql())
 
+        anon_unquoted_str = parse_one("anonymous(x, y)")
+        self.assertEqual(optimizer.simplify.gen(anon_unquoted_str), "ANONYMOUS x,y")
+
+        anon_unquoted_identifier = exp.Anonymous(
+            this=exp.to_identifier("anonymous"), expressions=[exp.column("x"), exp.column("y")]
+        )
+        self.assertEqual(optimizer.simplify.gen(anon_unquoted_identifier), "ANONYMOUS x,y")
+
+        anon_quoted = parse_one('"anonymous"(x, y)')
+        self.assertEqual(optimizer.simplify.gen(anon_quoted), '"anonymous" x,y')
+
+        with self.assertRaises(ValueError) as e:
+            anon_invalid = exp.Anonymous(this=5)
+            optimizer.simplify.gen(anon_invalid)
+
+        self.assertIn("Anonymous.this expects a str or an Identifier, got 'int'.", str(e.exception))
+
     def test_unnest_subqueries(self):
         self.check_file(
             "unnest_subqueries",
@@ -982,9 +1009,12 @@ FROM READ_CSV('tests/fixtures/optimizer/tpc-h/nation.csv.gz', 'delimiter', '|')
         self.assertEqual(expression.selects[0].type.sql(), "ARRAY<INT>")
 
         schema = MappingSchema({"t": {"c": "STRUCT<`f` STRING>"}}, dialect="bigquery")
-        expression = annotate_types(parse_one("SELECT t.c FROM t"), schema=schema)
+        expression = annotate_types(parse_one("SELECT t.c, [t.c] FROM t"), schema=schema)
 
         self.assertEqual(expression.selects[0].type.sql(dialect="bigquery"), "STRUCT<`f` STRING>")
+        self.assertEqual(
+            expression.selects[1].type.sql(dialect="bigquery"), "ARRAY<STRUCT<`f` STRING>>"
+        )
 
         expression = annotate_types(
             parse_one("SELECT unnest(t.x) FROM t AS t", dialect="postgres"),
@@ -1010,6 +1040,22 @@ FROM READ_CSV('tests/fixtures/optimizer/tpc-h/nation.csv.gz', 'delimiter', '|')
         self.assertEqual(exp.DataType.Type.USERDEFINED, expression.selects[0].type.this)
         self.assertEqual(expression.selects[0].type.sql(dialect="postgres"), "IPADDRESS")
 
+    def test_unnest_annotation(self):
+        expression = annotate_types(
+            optimizer.qualify.qualify(
+                parse_one(
+                    """
+                SELECT a, a.b, a.b.c FROM x, UNNEST(x.a) AS a
+                """,
+                    read="bigquery",
+                )
+            ),
+            schema={"x": {"a": "ARRAY<STRUCT<b STRUCT<c int>>>"}},
+        )
+        self.assertEqual(expression.selects[0].type, exp.DataType.build("STRUCT<b STRUCT<c int>>"))
+        self.assertEqual(expression.selects[1].type, exp.DataType.build("STRUCT<c int>"))
+        self.assertEqual(expression.selects[2].type, exp.DataType.build("int"))
+
     def test_recursive_cte(self):
         query = parse_one(
             """