feat: cache external models for faster loading (#4029)

Author: tobymao

sqlmesh/core/loader.py

@@ -23,7 +23,6 @@
 from sqlmesh.core.metric import Metric, MetricMeta, expand_metrics, load_metric_ddl
 from sqlmesh.core.model import (
     Model,
-    ExternalModel,
     ModelCache,
     SeedModel,
     create_external_model,
@@ -59,6 +58,14 @@ class LoadedProject:
     user_rules: RuleSet
 
 
+class CacheBase(abc.ABC):
+    @abc.abstractmethod
+    def get_or_load_models(
+        self, target_path: Path, loader: t.Callable[[], t.List[Model]]
+    ) -> t.List[Model]:
+        """Get or load all models from cache."""
+
+
 class Loader(abc.ABC):
     """Abstract base class to load macros and models for a context"""
 
@@ -192,6 +199,7 @@ def _load_metrics(self) -> UniqueKeyDict[str, MetricMeta]:
     def _load_external_models(
         self,
         audits: UniqueKeyDict[str, ModelAudit],
+        cache: CacheBase,
         gateway: t.Optional[str] = None,
     ) -> UniqueKeyDict[str, Model]:
         models: UniqueKeyDict[str, Model] = UniqueKeyDict("models")
@@ -208,32 +216,39 @@ def _load_external_models(
         if external_models_path.exists() and external_models_path.is_dir():
             paths_to_load.extend(self._glob_paths(external_models_path, extension=".yaml"))
 
+        def _load() -> t.List[Model]:
+            try:
+                with open(path, "r", encoding="utf-8") as file:
+                    return [
+                        create_external_model(
+                            defaults=self.config.model_defaults.dict(),
+                            path=path,
+                            project=self.config.project,
+                            audit_definitions=audits,
+                            **{
+                                "dialect": self.config.model_defaults.dialect,
+                                "default_catalog": self.context.default_catalog,
+                                **row,
+                            },
+                        )
+                        for row in YAML().load(file.read())
+                    ]
+            except Exception as ex:
+                raise ConfigError(f"Failed to load model definition at '{path}'.\n{ex}")
+
         for path in paths_to_load:
             self._track_file(path)
 
-            with open(path, "r", encoding="utf-8") as file:
-                external_models: t.List[ExternalModel] = []
-                for row in YAML().load(file.read()):
-                    model = create_external_model(
-                        defaults=self.config.model_defaults.dict(),
-                        path=path,
-                        project=self.config.project,
-                        audit_definitions=audits,
-                        **{
-                            "dialect": self.config.model_defaults.dialect,
-                            "default_catalog": self.context.default_catalog,
-                            **row,
-                        },
-                    )
-                    external_models.append(model)
-
-                # external models with no explicit gateway defined form the base set
-                for model in (e for e in external_models if e.gateway is None):
+            external_models = cache.get_or_load_models(path, _load)
+            # external models with no explicit gateway defined form the base set
+            for model in external_models:
+                if model.gateway is None:
                     models[model.fqn] = model
 
-                # however, if there is a gateway defined, gateway-specific models take precedence
-                if gateway:
-                    for model in (e for e in external_models if e.gateway == gateway):
+            # however, if there is a gateway defined, gateway-specific models take precedence
+            if gateway:
+                for model in external_models:
+                    if model.gateway == gateway:
                         models.update({model.fqn: model})
 
         return models
@@ -396,8 +411,9 @@ def _load_models(
         Loads all of the models within the model directory with their associated
         audits into a Dict and creates the dag
         """
-        sql_models = self._load_sql_models(macros, jinja_macros, audits, signals)
-        external_models = self._load_external_models(audits, gateway)
+        cache = SqlMeshLoader._Cache(self, self.config_path)
+        sql_models = self._load_sql_models(macros, jinja_macros, audits, signals, cache)
+        external_models = self._load_external_models(audits, cache, gateway)
         python_models = self._load_python_models(macros, jinja_macros, audits, signals)
 
         all_model_names = list(sql_models) + list(external_models) + list(python_models)
@@ -413,10 +429,10 @@ def _load_sql_models(
         jinja_macros: JinjaMacroRegistry,
         audits: UniqueKeyDict[str, ModelAudit],
         signals: UniqueKeyDict[str, signal],
+        cache: CacheBase,
     ) -> UniqueKeyDict[str, Model]:
         """Loads the sql models into a Dict"""
         models: UniqueKeyDict[str, Model] = UniqueKeyDict("models")
-        cache = SqlMeshLoader._Cache(self, self.config_path)
 
         for path in self._glob_paths(
             self.config_path / c.MODELS,
@@ -662,7 +678,7 @@ def _load_linting_rules(self) -> RuleSet:
 
         return RuleSet(user_rules.values())
 
-    class _Cache:
+    class _Cache(CacheBase):
         def __init__(self, loader: SqlMeshLoader, config_path: Path):
             self._loader = loader
             self.config_path = config_path

sqlmesh/core/model/cache.py

@@ -12,7 +12,7 @@
 from sqlglot.schema import MappingSchema
 
 from sqlmesh.core import constants as c
-from sqlmesh.core.model.definition import Model, SqlModel, _Model
+from sqlmesh.core.model.definition import ExternalModel, Model, SqlModel, _Model
 from sqlmesh.utils.cache import FileCache
 from sqlmesh.utils.hashing import crc32
 
@@ -59,7 +59,7 @@ def get_or_load(
             return cache_entry
 
         models = loader()
-        if isinstance(models, list) and isinstance(seq_get(models, 0), SqlModel):
+        if isinstance(models, list) and isinstance(seq_get(models, 0), (SqlModel, ExternalModel)):
             # make sure we preload full_depends_on
             for model in models:
                 model.full_depends_on

sqlmesh/dbt/loader.py

@@ -11,7 +11,7 @@
     GatewayConfig,
     ModelDefaultsConfig,
 )
-from sqlmesh.core.loader import LoadedProject, Loader
+from sqlmesh.core.loader import CacheBase, LoadedProject, Loader
 from sqlmesh.core.macros import MacroRegistry, macro
 from sqlmesh.core.model import Model, ModelCache
 from sqlmesh.core.signal import signal
@@ -145,7 +145,7 @@ def _to_sqlmesh(config: BMC, context: DbtContext) -> Model:
 
                     models[sqlmesh_model.fqn] = sqlmesh_model
 
-            models.update(self._load_external_models(audits))
+            models.update(self._load_external_models(audits, cache))
 
         return models
 
@@ -255,7 +255,7 @@ def _compute_yaml_max_mtime_per_subfolder(self, root: Path) -> t.Dict[Path, floa
 
         return result
 
-    class _Cache:
+    class _Cache(CacheBase):
         MAX_ENTRY_NAME_LENGTH = 200
 
         def __init__(

tests/core/test_loader.py

@@ -201,8 +201,3 @@ def my_model(context, **kwargs):
     assert model.description == "model_payload_a"
     path_b.write_text(model_payload_b)
     context.load()  # raise no error to duplicate key if the functions are identical (by registry class_method)
-    model = context.get_model(f"{model_name}")
-    assert (
-        model.description != "model_payload_b"
-    )  # model will not be overwritten by model_payload_b
-    assert model.description == "model_payload_a"