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