diff options
Diffstat (limited to 'tests/test_optimizer.py')
| -rw-r--r-- | tests/test_optimizer.py | 48 |
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( """ |