Fix!: normalize identifiers and model names (#930)

* Fix!: normalize identifiers and model names during optimization

* Revert query to query_ rename

* Typo

* Test fixups

* PR feedback

* Fix fingerprints

* Quote identifiers at the very end

* Ensure we use get_model, get_snapshot consistently, do some cleanups

* Make depends_on validator more robust by using the dialect

* Update test to reflect depends_on validator changes

* PR feedback

* Bump sqlglot

* Fix table diff test

Author: georgesittas

setup.py

@@ -44,7 +44,7 @@
         "requests",
         "rich",
         "ruamel.yaml",
-        "sqlglot~=14.1.0",
+        "sqlglot~=15.0.0",
         "fsspec",
     ],
     extras_require={

sqlmesh/cli/main.py

@@ -109,13 +109,8 @@ def render(
     dialect: t.Optional[str] = None,
 ) -> None:
     """Renders a model's query, optionally expanding referenced models."""
-    snapshot = ctx.obj.snapshots.get(model)
-
-    if not snapshot:
-        raise click.ClickException(f"Model `{model}` not found.")
-
     rendered = ctx.obj.render(
-        snapshot,
+        model,
         start=start,
         end=end,
         latest=latest,

sqlmesh/core/context.py

@@ -49,7 +49,12 @@
 from sqlmesh.core.config import Config, load_config_from_paths
 from sqlmesh.core.console import Console, get_console
 from sqlmesh.core.context_diff import ContextDiff
-from sqlmesh.core.dialect import format_model_expressions, pandas_to_sql, parse
+from sqlmesh.core.dialect import (
+    format_model_expressions,
+    normalize_model_name,
+    pandas_to_sql,
+    parse,
+)
 from sqlmesh.core.engine_adapter import EngineAdapter
 from sqlmesh.core.environment import Environment
 from sqlmesh.core.loader import Loader, SqlMeshLoader, update_model_schemas
@@ -277,13 +282,13 @@ def upsert_model(self, model: t.Union[str, Model], **kwargs: t.Any) -> Model:
         Returns:
             A new instance of the updated or inserted model.
         """
-        if isinstance(model, str):
-            model = self._models[model]
+        model = self.get_model(model, raise_if_missing=True)
+        path = model._path
 
-        path = model._path  # type: ignore
         # model.copy() can't be used here due to a cached state that can be a part of a model instance.
         model = t.cast(Model, type(model)(**{**t.cast(Model, model).dict(), **kwargs}))
         model._path = path
+
         self._models.update({model.name: model})
 
         self._add_model_to_dag(model)
@@ -409,14 +414,16 @@ def get_model(
             The expected model.
         """
         if isinstance(model_or_snapshot, str):
-            model = self._models.get(model_or_snapshot)
+            normalized_name = normalize_model_name(model_or_snapshot, dialect=self.config.dialect)
+            model = self._models.get(normalized_name)
         elif isinstance(model_or_snapshot, Snapshot):
             model = model_or_snapshot.model
         else:
             model = model_or_snapshot
 
         if raise_if_missing and not model:
             raise SQLMeshError(f"Cannot find model for '{model_or_snapshot}'")
+
         return model
 
     @t.overload
@@ -444,7 +451,8 @@ def get_snapshot(
             The expected snapshot.
         """
         if isinstance(model_or_snapshot, str):
-            snapshot = self.snapshots.get(model_or_snapshot)
+            normalized_name = normalize_model_name(model_or_snapshot, dialect=self.config.dialect)
+            snapshot = self.snapshots.get(normalized_name)
         elif isinstance(model_or_snapshot, Snapshot):
             snapshot = model_or_snapshot
         else:
@@ -908,7 +916,9 @@ def audit(
         """
 
         snapshots = (
-            [self.snapshots[model] for model in models] if models else self.snapshots.values()
+            [self.get_snapshot(model, raise_if_missing=True) for model in models]
+            if models
+            else self.snapshots.values()
         )
 
         num_audits = sum(len(snapshot.model.audits) for snapshot in snapshots)
@@ -942,6 +952,7 @@ def audit(
             )
             self.console.log_status_update(f"Got {error.count} results, expected 0.")
             self.console.show_sql(f"{error.query}")
+
         self.console.log_status_update("Done.")
 
     def migrate(self) -> None:

sqlmesh/core/dialect.py

@@ -8,6 +8,8 @@
 import pandas as pd
 from jinja2.meta import find_undeclared_variables
 from sqlglot import Dialect, Generator, Parser, TokenType, exp
+from sqlglot.dialects.dialect import DialectType
+from sqlglot.optimizer.normalize_identifiers import normalize_identifiers
 
 from sqlmesh.core.constants import MAX_MODEL_DEFINITION_SIZE
 from sqlmesh.utils.jinja import ENVIRONMENT
@@ -544,3 +546,9 @@ def pandas_to_sql(
         batch_size=batch_size,
         alias=alias,
     )
+
+
+def normalize_model_name(table: str | exp.Table, dialect: DialectType = None) -> str:
+    return exp.table_name(
+        normalize_identifiers(exp.to_table(table, dialect=dialect), dialect=dialect)
+    )

sqlmesh/core/loader.py

@@ -12,6 +12,7 @@
 from pathlib import Path
 
 from ruamel.yaml import YAML
+from sqlglot import exp
 from sqlglot.errors import SqlglotError
 from sqlglot.optimizer.qualify_columns import validate_qualify_columns
 from sqlglot.schema import MappingSchema, nested_set
@@ -39,7 +40,7 @@
 
 
 def update_model_schemas(dag: DAG[str], models: UniqueKeyDict[str, Model]) -> None:
-    schema = MappingSchema()
+    schema = MappingSchema(normalize=False)
 
     for name in dag.sorted():
         model = models.get(name)
@@ -52,9 +53,7 @@ def update_model_schemas(dag: DAG[str], models: UniqueKeyDict[str, Model]) -> No
 
         for dep in model.depends_on:
             external = external or dep not in models
-            table = schema._normalize_table(
-                schema._ensure_table(dep, dialect=model.dialect), dialect=model.dialect
-            )
+            table = exp.to_table(dep, dialect=model.dialect)
             mapping_schema = schema.find(table)
 
             if mapping_schema:

sqlmesh/core/model/definition.py

@@ -654,7 +654,7 @@ def columns_to_types(self) -> t.Dict[str, exp.DataType]:
         if self._columns_to_types is None:
             self._columns_to_types = {
                 expression.output_name: expression.type or exp.DataType.build("unknown")
-                for expression in self._query_renderer.render().expressions
+                for expression in self._query_renderer.render().selects
             }
 
         return self._columns_to_types
@@ -667,7 +667,7 @@ def column_descriptions(self) -> t.Dict[str, str]:
         if self._column_descriptions is None:
             self._column_descriptions = {
                 select.alias: "\n".join(comment.strip() for comment in select.comments)
-                for select in self.render_query().expressions
+                for select in self.render_query().selects
                 if select.comments
             }
         return self._column_descriptions
@@ -678,9 +678,7 @@ def validate_definition(self) -> None:
         if not isinstance(query, exp.Subqueryable):
             raise_config_error("Missing SELECT query in the model definition", self._path)
 
-        projection_list = (
-            query.expressions if not isinstance(query, exp.Union) else query.this.expressions
-        )
+        projection_list = query.selects
         if not projection_list:
             raise_config_error("Query missing select statements", self._path)
 

sqlmesh/core/model/meta.py

@@ -41,9 +41,9 @@ class IntervalUnit(str, Enum):
 class ModelMeta(PydanticModel):
     """Metadata for models which can be defined in SQL."""
 
+    dialect: str = ""
     name: str
     kind: ModelKind = ModelKind(name=ModelKindName.VIEW)
-    dialect: str = ""
     cron: str = "@daily"
     owner: t.Optional[str]
     description: t.Optional[str]
@@ -65,6 +65,10 @@ class ModelMeta(PydanticModel):
 
     _model_kind_validator = ModelKind.field_validator()
 
+    @validator("name", pre=True)
+    def _name_validator(cls, v: t.Any, values: t.Dict[str, t.Any]) -> str:
+        return d.normalize_model_name(v, dialect=values.get("dialect"))
+
     @validator("audits", pre=True)
     def _audits_validator(cls, v: t.Any) -> t.Any:
         def extract(v: exp.Expression) -> t.Tuple[str, t.Dict[str, str]]:
@@ -144,14 +148,19 @@ def _columns_validator(
         return v
 
     @validator("depends_on_", pre=True)
-    def _depends_on_validator(cls, v: t.Any) -> t.Optional[t.Set[str]]:
+    def _depends_on_validator(cls, v: t.Any, values: t.Dict[str, t.Any]) -> t.Optional[t.Set[str]]:
+        dialect = values.get("dialect")
+
         if isinstance(v, (exp.Array, exp.Tuple)):
             return {
-                exp.table_name(table.name if table.is_string else table.sql())
+                d.normalize_model_name(
+                    table.name if table.is_string else table.sql(dialect=dialect), dialect=dialect
+                )
                 for table in v.expressions
             }
         if isinstance(v, exp.Expression):
-            return {exp.table_name(v.sql())}
+            return {d.normalize_model_name(v.sql(dialect=dialect), dialect=dialect)}
+
         return v
 
     @validator("start", pre=True)

sqlmesh/core/renderer.py

@@ -7,9 +7,8 @@
 from sqlglot import exp, parse_one
 from sqlglot.errors import SqlglotError
 from sqlglot.optimizer.annotate_types import annotate_types
-from sqlglot.optimizer.lower_identities import lower_identities
-from sqlglot.optimizer.qualify_columns import qualify_columns
-from sqlglot.optimizer.qualify_tables import qualify_tables
+from sqlglot.optimizer.qualify import qualify
+from sqlglot.optimizer.qualify_columns import quote_identifiers
 from sqlglot.optimizer.simplify import simplify
 from sqlglot.schema import ensure_schema
 
@@ -259,7 +258,7 @@ def render(
         if not isinstance(query, exp.Subqueryable):
             raise_config_error(f"Query needs to be a SELECT or a UNION {query}.", self._path)
 
-        return t.cast(exp.Subqueryable, query)
+        return t.cast(exp.Subqueryable, quote_identifiers(query, dialect=self._dialect))
 
     def time_column_filter(self, start: TimeLike, end: TimeLike) -> exp.Between:
         """Returns a between statement with the properly formatted time column."""
@@ -273,19 +272,25 @@ def time_column_filter(self, start: TimeLike, end: TimeLike) -> exp.Between:
         )
 
     def _optimize_query(self, query: exp.Expression) -> exp.Expression:
-        schema = ensure_schema(self.schema, dialect=self._dialect)
-
+        # We don't want to normalize names in the schema because that's handled by the optimizer
+        schema = ensure_schema(self.schema, dialect=self._dialect, normalize=False)
         query = t.cast(exp.Subqueryable, query.copy())
-        lower_identities(query)
-        qualify_tables(query)
 
         try:
+            qualify(
+                query,
+                dialect=self._dialect,
+                schema=schema,
+                infer_schema=False,
+                validate_qualify_columns=False,
+                qualify_columns=not schema.empty,
+                quote_identifiers=False,
+            )
+
             if schema.empty:
                 for select in query.selects:
                     if not isinstance(select, exp.Alias) and select.output_name not in ("*", ""):
                         select.replace(exp.alias_(select, select.output_name))
-            else:
-                qualify_columns(query, schema=schema, infer_schema=False)
         except SqlglotError as ex:
             raise_config_error(
                 f"Error qualifying columns, the column may not exist or is ambiguous. {ex}",

sqlmesh/magics.py

@@ -14,7 +14,7 @@
 from sqlmesh.core.dialect import format_model_expressions, parse
 from sqlmesh.core.model import load_model
 from sqlmesh.core.test import ModelTestMetadata, get_all_model_tests
-from sqlmesh.utils.errors import MagicError, MissingContextException, SQLMeshError
+from sqlmesh.utils.errors import MagicError, MissingContextException
 from sqlmesh.utils.yaml import dumps as yaml_dumps
 from sqlmesh.utils.yaml import load as yaml_load
 
@@ -66,10 +66,7 @@ def context(self, line: str) -> None:
     def model(self, line: str, sql: t.Optional[str] = None) -> None:
         """Renders the model and automatically fills in an editable cell with the model definition."""
         args = parse_argstring(self.model, line)
-        model = self._context.get_model(t.cast(str, args.model))
-
-        if not model:
-            raise SQLMeshError(f"Cannot find {model}")
+        model = self._context.get_model(args.model, raise_if_missing=True)
 
         if sql:
             config = self._context.config_for_model(model)
@@ -142,13 +139,16 @@ def test(self, line: str, test_def_raw: t.Optional[str] = None) -> None:
                     f"Test found that does not have `model` defined: {model_test_metadata.path}"
                 )
             tests[model][model_test_metadata.test_name] = model_test_metadata
+
+        model = self._context.get_model(args.model, raise_if_missing=True)
+
         if args.ls:
             # TODO: Provide better UI for displaying tests
-            for test_name in tests[args.model]:
+            for test_name in tests[model.name]:
                 self._context.console.log_status_update(test_name)
             return
 
-        test = tests[args.model][args.test_name]
+        test = tests[model.name][args.test_name]
         test_def = yaml_load(test_def_raw) if test_def_raw else test.body
         test_def_output = yaml_dumps(test_def)
 

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" <= '1970-01-01' AND "ds" >= '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" <= '1970-01-01' AND "ds" >= '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" <= '1970-01-01' AND "ds" >= '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" <= '1970-01-01' AND "ds" >= '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" <= '1970-01-01' AND "ds" >= '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" <= '1970-01-01' AND "ds" >= '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" <= '1970-01-01' AND "ds" >= '1970-01-01') AS "_q_0" HAVING COUNT(*) <= 0 LIMIT 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" <= '1970-01-01' AND "ds" >= '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" <= '1970-01-01' AND "ds" >= '1970-01-01') AS "_q_0" WHERE NOT "a" >= "b" OR NOT "c" + "d" - "e" < 1.0"""
     )