Fix: use the models' dialects in schema methods (#889)

* Fix: use the models' dialects in schema methods

* Fix sqlglot version

Author: georgesittas

setup.py

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

sqlmesh/core/dialect.py

@@ -423,7 +423,7 @@ def parse(sql: str, default_dialect: str | None = None) -> t.List[exp.Expression
     for chunk, is_jinja in chunks:
         if is_jinja:
             start, *_, end = chunk
-            segment = sql[start.start : end.end + 1]
+            segment = sql[start.start : end.end + 2]
             variables = [
                 exp.Literal.string(var)
                 for var in find_undeclared_variables(ENVIRONMENT.parse(segment))

sqlmesh/core/loader.py

@@ -54,7 +54,9 @@ 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))
+            table = schema._normalize_table(
+                schema._ensure_table(dep, dialect=model.dialect), dialect=model.dialect
+            )
             mapping_schema = schema.find(table)
 
             if mapping_schema:
@@ -64,7 +66,7 @@ def update_model_schemas(dag: DAG[str], models: UniqueKeyDict[str, Model]) -> No
                     {k: str(v) for k, v in mapping_schema.items()},
                 )
 
-        schema.add_table(name, model.columns_to_types)
+        schema.add_table(name, model.columns_to_types, dialect=model.dialect)
 
         if external:
             if "*" in model.columns_to_types:
@@ -129,8 +131,10 @@ def load(self, context: Context, update_schemas: bool = True) -> LoadedProject:
 
         macros, hooks, jinja_macros = self._load_scripts()
         models = self._load_models(macros, hooks, jinja_macros)
+
         for model in models.values():
             self._add_model_to_dag(model)
+
         if update_schemas:
             update_model_schemas(self._dag, models)
 

sqlmesh/core/renderer.py

@@ -278,10 +278,12 @@ def update_cache(
         self._query_cache[_dates(start, end, latest)] = query
 
     def _optimize_query(self, query: exp.Expression) -> exp.Expression:
-        schema = ensure_schema(self.schema)
+        schema = ensure_schema(self.schema, dialect=self._dialect)
+
         query = query.copy()
         lower_identities(query)
         qualify_tables(query)
+
         try:
             if not schema.empty:
                 qualify_columns(query, schema=schema, infer_schema=False)
@@ -290,4 +292,5 @@ def _optimize_query(self, query: exp.Expression) -> exp.Expression:
                 f"Error qualifying columns, the column may not exist or is ambiguous. {ex}",
                 self._path,
             )
+
         return annotate_types(simplify(query), schema=schema)

tests/core/test_model.py

@@ -1072,12 +1072,13 @@ def test_model_cache(tmp_path: Path, mocker: MockerFixture):
 def test_model_ctas_query():
     expressions = parse(
         """
-        MODEL (name db.table, kind FULL);
+        MODEL (name `a-b-c.table`, kind FULL, dialect bigquery);
         SELECT 1 as a
-        """
+        """,
+        read="bigquery",
     )
 
-    assert load_model(expressions).ctas_query({}).sql() == "SELECT 1 AS a"
+    assert load_model(expressions, dialect="bigquery").ctas_query({}).sql() == "SELECT 1 AS a"
 
     expressions = parse(
         """