Fix!: tables with dots and more conservative optimization (#1052)

Author: tobymao

setup.py

@@ -45,7 +45,7 @@
         "requests",
         "rich",
         "ruamel.yaml",
-        "sqlglot~=16.4.2",
+        "sqlglot~=16.7.0",
         "fsspec",
     ],
     extras_require={

sqlmesh/core/dialect.py

@@ -673,5 +673,5 @@ def find_tables(expression: exp.Expression, dialect: DialectType = None) -> t.Se
         normalize_model_name(table, dialect=dialect)
         for scope in traverse_scope(expression)
         for table in scope.tables
-        if isinstance(table.this, exp.Identifier) and exp.table_name(table) not in scope.cte_sources
+        if not isinstance(table.this, exp.Func) and exp.table_name(table) not in scope.cte_sources
     }

sqlmesh/core/loader.py

@@ -12,7 +12,6 @@
 from pathlib import Path
 
 from sqlglot.errors import SqlglotError
-from sqlglot.optimizer.qualify_columns import validate_qualify_columns
 from sqlglot.schema import MappingSchema
 
 from sqlmesh.core import constants as c
@@ -24,7 +23,6 @@
     ModelCache,
     OptimizedQueryCache,
     SeedModel,
-    SqlModel,
     create_external_model,
     load_model,
 )
@@ -54,23 +52,12 @@ def update_model_schemas(
             continue
 
         model.update_schema(schema)
-
-        cache_hit = optimized_query_cache.with_optimized_query(model)
+        optimized_query_cache.with_optimized_query(model)
 
         columns_to_types = model.columns_to_types
         if columns_to_types is not None:
             schema.add_table(name, columns_to_types, dialect=model.dialect)
 
-        if isinstance(model, SqlModel) and model.mapping_schema and not cache_hit:
-            query = model.render_query()
-            if query is not None:
-                try:
-                    validate_qualify_columns(query)
-                except SqlglotError as e:
-                    raise ConfigError(
-                        f"Column references could not be resolved for model '{name}' at '{model._path}'. {e}"
-                    )
-
 
 @dataclass
 class LoadedProject:

sqlmesh/core/renderer.py

@@ -148,7 +148,9 @@ def render(
             except MacroEvalError as ex:
                 raise_config_error(f"Failed to resolve macro for expression. {ex}", self._path)
 
-            _normalize_and_quote(expression, self._dialect)
+            if expression:
+                with _normalize_and_quote(expression, self._dialect) as expression:
+                    pass
 
             self._cache[cache_key] = expression
 
@@ -383,12 +385,13 @@ def _optimize_query(self, query: exp.Subqueryable) -> exp.Subqueryable:
             if should_optimize:
                 query = query.copy()
 
-                qualify(
-                    query,
-                    dialect=self._dialect,
-                    schema=schema,
-                    infer_schema=False,
-                    validate_qualify_columns=False,
+                simplify(
+                    qualify(
+                        query,
+                        dialect=self._dialect,
+                        schema=schema,
+                        infer_schema=False,
+                    )
                 )
         except SqlglotError as ex:
             failure = True
@@ -405,13 +408,13 @@ def _optimize_query(self, query: exp.Subqueryable) -> exp.Subqueryable:
                         ):
                             select.replace(exp.alias_(select, select.output_name))
 
-        return annotate_types(simplify(query), schema=schema)
+        return annotate_types(query, schema=schema)
 
 
 @contextmanager
 def _normalize_and_quote(query: E, dialect: str) -> t.Iterator[E]:
-    normalize_identifiers(query, dialect=dialect)
     qualify_tables(query)
+    normalize_identifiers(query, dialect=dialect)
     yield query
     quote_identifiers(query, dialect=dialect)
 

tests/core/test_audit.py

@@ -144,7 +144,7 @@ def test_no_query():
 
 
 def test_macro(model: Model):
-    expected_query = """SELECT * FROM (SELECT * FROM "db"."test_model" AS "test_model" WHERE "ds" <= '1970-01-01' AND "ds" >= '1970-01-01') AS "_q_0" WHERE "a" IS NULL"""
+    expected_query = """SELECT * FROM (SELECT * FROM "db"."test_model" AS "test_model" WHERE "ds" BETWEEN '1970-01-01' AND '1970-01-01') AS "_q_0" WHERE "a" IS NULL"""
 
     audit = Audit(
         name="test_audit",
@@ -167,7 +167,7 @@ def test_not_null_audit(model: Model):
     )
     assert (
         rendered_query_a.sql()
-        == """SELECT * FROM (SELECT * FROM "db"."test_model" AS "test_model" WHERE "ds" <= '1970-01-01' AND "ds" >= '1970-01-01') AS "_q_0" WHERE "a" IS NULL"""
+        == """SELECT * FROM (SELECT * FROM "db"."test_model" AS "test_model" WHERE "ds" BETWEEN '1970-01-01' AND '1970-01-01') AS "_q_0" WHERE "a" IS NULL"""
     )
 
     rendered_query_a_and_b = builtin.not_null_audit.render_query(
@@ -176,23 +176,23 @@ def test_not_null_audit(model: Model):
     )
     assert (
         rendered_query_a_and_b.sql()
-        == """SELECT * FROM (SELECT * FROM "db"."test_model" AS "test_model" WHERE "ds" <= '1970-01-01' AND "ds" >= '1970-01-01') AS "_q_0" WHERE "a" IS NULL OR "b" IS NULL"""
+        == """SELECT * FROM (SELECT * FROM "db"."test_model" AS "test_model" WHERE "ds" BETWEEN '1970-01-01' AND '1970-01-01') AS "_q_0" WHERE "a" IS NULL OR "b" IS NULL"""
     )
 
 
 def test_unique_values_audit(model: Model):
     rendered_query_a = builtin.unique_values_audit.render_query(model, columns=[exp.to_column("a")])
     assert (
         rendered_query_a.sql()
-        == """SELECT * FROM (SELECT ROW_NUMBER() OVER (PARTITION BY "a" ORDER BY 1) AS "a_rank" FROM (SELECT * FROM "db"."test_model" AS "test_model" WHERE "ds" <= '1970-01-01' AND "ds" >= '1970-01-01') AS "_q_0") AS "_q_1" WHERE "a_rank" > 1"""
+        == """SELECT * FROM (SELECT ROW_NUMBER() OVER (PARTITION BY "a" ORDER BY 1) AS "a_rank" FROM (SELECT * FROM "db"."test_model" AS "test_model" WHERE "ds" BETWEEN '1970-01-01' AND '1970-01-01') AS "_q_0") AS "_q_1" WHERE "a_rank" > 1"""
     )
 
     rendered_query_a_and_b = builtin.unique_values_audit.render_query(
         model, columns=[exp.to_column("a"), exp.to_column("b")]
     )
     assert (
         rendered_query_a_and_b.sql()
-        == """SELECT * FROM (SELECT ROW_NUMBER() OVER (PARTITION BY "a" ORDER BY 1) AS "a_rank", ROW_NUMBER() OVER (PARTITION BY "b" ORDER BY 1) AS "b_rank" FROM (SELECT * FROM "db"."test_model" AS "test_model" WHERE "ds" <= '1970-01-01' AND "ds" >= '1970-01-01') AS "_q_0") AS "_q_1" WHERE "a_rank" > 1 OR "b_rank" > 1"""
+        == """SELECT * FROM (SELECT ROW_NUMBER() OVER (PARTITION BY "a" ORDER BY 1) AS "a_rank", ROW_NUMBER() OVER (PARTITION BY "b" ORDER BY 1) AS "b_rank" FROM (SELECT * FROM "db"."test_model" AS "test_model" WHERE "ds" BETWEEN '1970-01-01' AND '1970-01-01') AS "_q_0") AS "_q_1" WHERE "a_rank" > 1 OR "b_rank" > 1"""
     )
 
 
@@ -204,7 +204,7 @@ def test_accepted_values_audit(model: Model):
     )
     assert (
         rendered_query.sql()
-        == """SELECT * FROM (SELECT * FROM "db"."test_model" AS "test_model" WHERE "ds" <= '1970-01-01' AND "ds" >= '1970-01-01') AS "_q_0" WHERE NOT "a" IN ('value_a', 'value_b')"""
+        == """SELECT * FROM (SELECT * FROM "db"."test_model" AS "test_model" WHERE "ds" BETWEEN '1970-01-01' AND '1970-01-01') AS "_q_0" WHERE NOT "a" IN ('value_a', 'value_b')"""
     )
 
 
@@ -215,7 +215,7 @@ def test_number_of_rows_audit(model: Model):
     )
     assert (
         rendered_query.sql()
-        == """SELECT 1 AS "1" FROM (SELECT * FROM "db"."test_model" AS "test_model" WHERE "ds" <= '1970-01-01' AND "ds" >= '1970-01-01') AS "_q_0" HAVING COUNT(*) <= 0 LIMIT 1"""
+        == """SELECT 1 AS "1" FROM (SELECT * FROM "db"."test_model" AS "test_model" WHERE "ds" BETWEEN '1970-01-01' AND '1970-01-01') AS "_q_0" HAVING COUNT(*) <= 0 LIMIT 0 + 1"""
     )
 
 
@@ -226,7 +226,7 @@ def test_forall_audit(model: Model):
     )
     assert (
         rendered_query_a.sql()
-        == '''SELECT * FROM (SELECT * FROM "db"."test_model" AS "test_model" WHERE "ds" <= '1970-01-01' AND "ds" >= '1970-01-01') AS "_q_0" WHERE NOT "a" >= "b"'''
+        == """SELECT * FROM (SELECT * FROM "db"."test_model" AS "test_model" WHERE "ds" BETWEEN '1970-01-01' AND '1970-01-01') AS "_q_0" WHERE NOT ("a" >= "b")"""
     )
 
     rendered_query_a = builtin.forall_audit.render_query(
@@ -235,5 +235,5 @@ def test_forall_audit(model: Model):
     )
     assert (
         rendered_query_a.sql()
-        == """SELECT * FROM (SELECT * FROM "db"."test_model" AS "test_model" WHERE "ds" <= '1970-01-01' AND "ds" >= '1970-01-01') AS "_q_0" WHERE NOT "a" >= "b" OR NOT "c" + "d" - "e" < 1.0"""
+        == """SELECT * FROM (SELECT * FROM "db"."test_model" AS "test_model" WHERE "ds" BETWEEN '1970-01-01' AND '1970-01-01') AS "_q_0" WHERE NOT ("a" >= "b") OR NOT ("c" + "d" - "e" < 1.0)"""
     )

tests/core/test_model.py

@@ -1,3 +1,4 @@
+import logging
 from datetime import date
 from pathlib import Path
 from unittest.mock import patch
@@ -195,6 +196,27 @@ def test_model_validation_union_query():
         model.validate_definition()
 
 
+def test_model_qualification():
+    logger = logging.getLogger("sqlmesh.core.renderer")
+    with patch.object(logger, "error") as mock_logger:
+        expressions = d.parse(
+            """
+            MODEL (
+                name db.table,
+                kind FULL,
+            );
+
+            SELECT a
+            """
+        )
+
+        model = load_model(expressions)
+        model.render_query(optimize=True)
+        assert (
+            mock_logger.call_args[0][0] == "%s for '%s', the column may not exist or is ambiguous"
+        )
+
+
 @pytest.mark.parametrize(
     "partition_by_input, partition_by_output, expected_exception",
     [
@@ -877,10 +899,8 @@ def test_render_query(assert_exp_eq):
           y AS y
         FROM x AS x
         WHERE
-          y <= '2020-10-28'
-          AND y <= DATE_STR_TO_DATE('2020-10-28')
-          AND y >= '2020-10-28'
-          AND y >= DATE_STR_TO_DATE('2020-10-28')
+          y BETWEEN DATE_STR_TO_DATE('2020-10-28') AND DATE_STR_TO_DATE('2020-10-28')
+          AND y BETWEEN '2020-10-28' AND '2020-10-28'
         """,
     )
     assert_exp_eq(
@@ -895,10 +915,8 @@ def test_render_query(assert_exp_eq):
             y AS y
           FROM x AS x
           WHERE
-            y <= '2020-10-28'
-            AND y <= DATE_STR_TO_DATE('2020-10-28')
-            AND y >= '2020-10-28'
-            AND y >= DATE_STR_TO_DATE('2020-10-28')
+            y BETWEEN DATE_STR_TO_DATE('2020-10-28') AND DATE_STR_TO_DATE('2020-10-28')
+            AND y BETWEEN '2020-10-28' AND '2020-10-28'
         ) AS _subquery
         WHERE
           y BETWEEN TIME_STR_TO_TIME('2020-10-28T00:00:00+00:00') AND TIME_STR_TO_TIME('2020-10-28T23:59:59.999000+00:00')
@@ -1190,7 +1208,7 @@ def test_parse(assert_exp_eq):
         ds AS ds
       FROM x AS x
       WHERE
-        ds <= '1970-01-01' AND ds >= '1970-01-01'
+        ds BETWEEN '1970-01-01' AND '1970-01-01'
     """,
     )
 

tests/dbt/test_transformation.py

@@ -228,7 +228,7 @@ def test_config_jinja(sushi_test_project: Project):
     context = sushi_test_project.context
     model = t.cast(SqlModel, model_config.to_sqlmesh(context))
     assert hook in model.pre_statements[0].sql()
-    assert model.render_pre_statements()[0].sql() == "bar"
+    assert model.render_pre_statements()[0].sql() == '"bar"'
 
 
 def test_this(assert_exp_eq, sushi_test_project: Project):