Fix: Check all dialects used in a project when retrieving a model by its name (#2801)

Author: izeigerman

sqlmesh/core/context.py

@@ -325,6 +325,8 @@ def __init__(
 
         self.path, self.config = t.cast(t.Tuple[Path, C], next(iter(self.configs.items())))
 
+        self._all_dialects: t.Set[str] = {self.config.dialect or ""}
+
         # This allows overriding the default dialect's normalization strategy, so for example
         # one can do `dialect="duckdb,normalization_strategy=lowercase"` and this will be
         # applied to the DuckDB dialect globally
@@ -435,6 +437,9 @@ def upsert_model(self, model: t.Union[str, Model], **kwargs: t.Any) -> Model:
             self.path,
         )
 
+        if model.dialect:
+            self._all_dialects.add(model.dialect)
+
         model.validate_definition()
 
         return model
@@ -517,6 +522,10 @@ def load(self, update_schemas: bool = True) -> GenericContext[C]:
                     f"Models and Standalone audits cannot have the same name: {duplicates}"
                 )
 
+            self._all_dialects = {m.dialect for m in self._models.values() if m.dialect} | {
+                self.default_dialect or ""
+            }
+
         analytics.collector.on_project_loaded(
             project_type=(
                 c.DBT if type(self._loader).__name__.lower().startswith(c.DBT) else c.NATIVE
@@ -624,21 +633,24 @@ def get_model(
             The expected model.
         """
         if isinstance(model_or_snapshot, str):
-            normalized_name = normalize_model_name(
-                model_or_snapshot,
-                dialect=self.default_dialect,
-                default_catalog=self.default_catalog,
-            )
-            model = self._models.get(normalized_name)
+            # We should try all dialects referenced in the project for cases when models use mixed dialects.
+            for dialect in self._all_dialects:
+                normalized_name = normalize_model_name(
+                    model_or_snapshot,
+                    dialect=dialect,
+                    default_catalog=self.default_catalog,
+                )
+                if normalized_name in self._models:
+                    return self._models[normalized_name]
         elif isinstance(model_or_snapshot, Snapshot):
-            model = model_or_snapshot.model
+            return model_or_snapshot.model
         else:
-            model = model_or_snapshot
+            return model_or_snapshot
 
-        if raise_if_missing and not model:
+        if raise_if_missing:
             raise SQLMeshError(f"Cannot find model for '{model_or_snapshot}'")
 
-        return model
+        return None
 
     @t.overload
     def get_snapshot(self, node_or_snapshot: NodeOrSnapshot) -> t.Optional[Snapshot]: ...

tests/core/test_context.py

@@ -765,6 +765,25 @@ def test_disabled_model(copy_to_temp_path):
     assert not context.get_model("sushi.disabled")
 
 
+def test_get_model_mixed_dialects(copy_to_temp_path):
+    path = copy_to_temp_path("examples/sushi")
+
+    context = Context(paths=path)
+    expression = d.parse(
+        """
+    MODEL(
+        name sushi.snowflake_dialect,
+        dialect snowflake,
+    );
+
+    SELECT 1"""
+    )
+    model = load_sql_based_model(expression, default_catalog=context.default_catalog)
+    context.upsert_model(model)
+
+    assert context.get_model("sushi.snowflake_dialect") == model
+
+
 def test_override_dialect_normalization_strategy():
     config = Config(
         model_defaults=ModelDefaultsConfig(dialect="duckdb,normalization_strategy=lowercase")