feat: cache external models and others for faster loading

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,
@@ -194,49 +193,7 @@ def _load_external_models(
         audits: UniqueKeyDict[str, ModelAudit],
         gateway: t.Optional[str] = None,
     ) -> UniqueKeyDict[str, Model]:
-        models: UniqueKeyDict[str, Model] = UniqueKeyDict("models")
-        external_models_yaml = Path(self.config_path / c.EXTERNAL_MODELS_YAML)
-        deprecated_yaml = Path(self.config_path / c.EXTERNAL_MODELS_DEPRECATED_YAML)
-        external_models_path = self.config_path / c.EXTERNAL_MODELS
-
-        paths_to_load = []
-        if external_models_yaml.exists():
-            paths_to_load.append(external_models_yaml)
-        elif deprecated_yaml.exists():
-            paths_to_load.append(deprecated_yaml)
-
-        if external_models_path.exists() and external_models_path.is_dir():
-            paths_to_load.extend(self._glob_paths(external_models_path, extension=".yaml"))
-
-        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):
-                    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):
-                        models.update({model.fqn: model})
-
-        return models
+        return UniqueKeyDict("external_models")
 
     def _load_requirements(self) -> t.Tuple[t.Dict[str, str], t.Set[str]]:
         """Loads Python dependencies from the lock file.
@@ -620,6 +577,59 @@ def _load_metrics(self) -> UniqueKeyDict[str, MetricMeta]:
 
         return metrics
 
+    def _load_external_models(
+        self,
+        audits: UniqueKeyDict[str, ModelAudit],
+        gateway: t.Optional[str] = None,
+    ) -> UniqueKeyDict[str, Model]:
+        models: UniqueKeyDict[str, Model] = UniqueKeyDict("models")
+        external_models_yaml = Path(self.config_path / c.EXTERNAL_MODELS_YAML)
+        deprecated_yaml = Path(self.config_path / c.EXTERNAL_MODELS_DEPRECATED_YAML)
+        external_models_path = self.config_path / c.EXTERNAL_MODELS
+
+        paths_to_load = []
+        if external_models_yaml.exists():
+            paths_to_load.append(external_models_yaml)
+        elif deprecated_yaml.exists():
+            paths_to_load.append(deprecated_yaml)
+
+        if external_models_path.exists() and external_models_path.is_dir():
+            paths_to_load.extend(self._glob_paths(external_models_path, extension=".yaml"))
+
+        cache = SqlMeshLoader._Cache(self, self.config_path)
+
+        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)
+
+            for model in cache.get_or_load_models(path, _load):
+                if model.gateway == gateway:
+                    models.update({model.fqn: model})
+                # external models with no explicit gateway defined form the base set
+                elif gateway and model.gateway is None and model.fqn not in models:
+                    models[model.fqn] = model
+
+        return models
+
     def _load_environment_statements(self, macros: MacroRegistry) -> EnvironmentStatements | None:
         """Loads environment statements."""
 

sqlmesh/core/model/cache.py

@@ -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):
             # make sure we preload full_depends_on
             for model in models:
                 model.full_depends_on