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,
@@ -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"""
 
@@ -69,6 +76,10 @@ def __init__(self, context: GenericContext, path: Path) -> None:
         self.config = self.context.configs[self.config_path]
         self._variables_by_gateway: t.Dict[str, t.Dict[str, t.Any]] = {}
 
+    @abc.abstractmethod
+    def _cache(self, **kwargs: t.Any) -> CacheBase:
+        """Returns an instance of a CacheBase."""
+
     def load(self) -> LoadedProject:
         """
         Loads all macros and models in the context's path.
@@ -208,32 +219,41 @@ 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}")
+
+        cache = self._cache()
+
         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
@@ -339,6 +359,9 @@ def _get_variables(self, gateway_name: t.Optional[str] = None) -> t.Dict[str, t.
 class SqlMeshLoader(Loader):
     """Loads macros and models for a context using the SQLMesh file formats"""
 
+    def _cache(self, **kwargs: t.Any) -> CacheBase:
+        return SqlMeshLoader._Cache(self, self.config_path)
+
     def _load_scripts(self) -> t.Tuple[MacroRegistry, JinjaMacroRegistry]:
         """Loads all user defined macros."""
         # Store a copy of the macro registry
@@ -416,7 +439,8 @@ def _load_sql_models(
     ) -> UniqueKeyDict[str, Model]:
         """Loads the sql models into a Dict"""
         models: UniqueKeyDict[str, Model] = UniqueKeyDict("models")
-        cache = SqlMeshLoader._Cache(self, self.config_path)
+
+        cache = self._cache()
 
         for path in self._glob_paths(
             self.config_path / c.MODELS,
@@ -662,7 +686,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

@@ -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

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
@@ -83,6 +83,9 @@ def __init__(self, context: GenericContext, path: Path) -> None:
         self._macros_max_mtime: t.Optional[float] = None
         super().__init__(context, path)
 
+    def _cache(self, **kwargs: t.Any) -> CacheBase:
+        return DbtLoader._Cache(self, **kwargs)
+
     def load(self) -> LoadedProject:
         self._projects = []
         return super().load()
@@ -120,7 +123,11 @@ def _to_sqlmesh(config: BMC, context: DbtContext) -> Model:
             yaml_max_mtimes = self._compute_yaml_max_mtime_per_subfolder(
                 project.context.project_root
             )
-            cache = DbtLoader._Cache(self, project, macros_max_mtime, yaml_max_mtimes)
+            cache = self._cache(
+                project=project,
+                macros_max_mtime=macros_max_mtime,
+                yaml_max_mtimes=yaml_max_mtimes,
+            )
 
             logger.debug("Converting models to sqlmesh")
             # Now that config is rendered, create the sqlmesh models
@@ -255,7 +262,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"