Merge branch 'main' into selective-snapshot-hydration

Author: cmgoffena13

sqlmesh/core/model/cache.py

@@ -112,22 +112,48 @@ def with_optimized_query(self, model: Model, name: t.Optional[str] = None) -> bo
         cache_entry = self._file_cache.get(name)
         if cache_entry:
             try:
-                # If the optimized rendered query is None, then there are likely adapter calls in the query
-                # that prevent us from rendering it at load time. This means that we can safely set the
-                # unoptimized cache to None as well to prevent attempts to render it downstream.
-                optimized = cache_entry.optimized_rendered_query is not None
-                model._query_renderer.update_cache(
-                    cache_entry.optimized_rendered_query,
-                    cache_entry.renderer_violations,
-                    optimized=optimized,
-                )
+                self.with_optimized_query_entry(model, cache_entry)
                 return True
             except Exception as ex:
                 logger.warning("Failed to load a cache entry '%s': %s", name, ex)
 
         self._put(name, model)
         return False
 
+    @staticmethod
+    def with_optimized_query_entry(model: Model, cache_entry: OptimizedQueryCacheEntry) -> None:
+        """Adds an already-decoded optimized query cache entry to a model."""
+        if not isinstance(model, SqlModel):
+            return
+
+        # If the optimized rendered query is None, then there are likely adapter calls in the
+        # query that prevent us from rendering it at load time. This means that we can safely
+        # set the unoptimized cache to None as well to prevent attempts to render it downstream.
+        optimized = cache_entry.optimized_rendered_query is not None
+        model._query_renderer.update_cache(
+            cache_entry.optimized_rendered_query,
+            cache_entry.renderer_violations,
+            optimized=optimized,
+        )
+
+    def get_or_create_entry(
+        self, model: Model, name: t.Optional[str] = None
+    ) -> t.Optional[OptimizedQueryCacheEntry]:
+        """Returns one decoded cache entry, creating it when necessary."""
+        if not isinstance(model, SqlModel):
+            return None
+
+        name = self._entry_name(model) if name is None else name
+        cache_entry = self._file_cache.get(name)
+        if cache_entry:
+            try:
+                self.with_optimized_query_entry(model, cache_entry)
+                return cache_entry
+            except Exception as ex:
+                logger.warning("Failed to load a cache entry '%s': %s", name, ex)
+
+        return self._put(name, model)
+
     def put(self, model: Model) -> t.Optional[str]:
         if not isinstance(model, SqlModel):
             return None
@@ -140,14 +166,15 @@ def put(self, model: Model) -> t.Optional[str]:
         self._put(name, model)
         return name
 
-    def _put(self, name: str, model: SqlModel) -> None:
+    def _put(self, name: str, model: SqlModel) -> OptimizedQueryCacheEntry:
         optimized_query = model.render_query()
 
         new_entry = OptimizedQueryCacheEntry(
             optimized_rendered_query=optimized_query,
             renderer_violations=model.violated_rules_for_query,
         )
         self._file_cache.put(name, value=new_entry)
+        return new_entry
 
     @staticmethod
     def _entry_name(model: SqlModel) -> str:
@@ -197,7 +224,13 @@ def load_optimized_query(
 
 def load_optimized_query_and_mapping(
     model: Model, mapping: t.Dict
-) -> t.Tuple[str, t.Optional[str], str, str, t.Dict]:
+) -> t.Tuple[
+    str,
+    t.Optional[OptimizedQueryCacheEntry],
+    str,
+    str,
+    t.Dict,
+]:
     assert _optimized_query_cache
 
     schema = MappingSchema(normalize=False)
@@ -207,13 +240,13 @@ def load_optimized_query_and_mapping(
 
     if isinstance(model, SqlModel):
         entry_name = _optimized_query_cache._entry_name(model)
-        _optimized_query_cache.with_optimized_query(model, entry_name)
+        cache_entry = _optimized_query_cache.get_or_create_entry(model, entry_name)
     else:
-        entry_name = None
+        cache_entry = None
 
     return (
         model.fqn,
-        entry_name,
+        cache_entry,
         model.data_hash,
         model.metadata_hash,
         model.mapping_schema,

sqlmesh/core/model/schema.py

@@ -87,13 +87,16 @@ def process_models(completed_model: t.Optional[Model] = None) -> None:
             for future in as_completed(futures):
                 try:
                     futures.remove(future)
-                    fqn, entry_name, data_hash, metadata_hash, mapping_schema = future.result()
+                    fqn, cache_entry, data_hash, metadata_hash, mapping_schema = future.result()
                     model = models[fqn]
                     model._data_hash = data_hash
                     model._metadata_hash = metadata_hash
                     if model.mapping_schema != mapping_schema:
                         model.set_mapping_schema(mapping_schema)
-                    optimized_query_cache.with_optimized_query(model, entry_name)
+                    if cache_entry is not None:
+                        # The worker has already decompressed and decoded this entry. Passing
+                        # it back avoids a second gzip/pickle read in the parent process.
+                        optimized_query_cache.with_optimized_query_entry(model, cache_entry)
                     _update_schema_with_model(schema, model)
                     process_models(completed_model=model)
                 except Exception as ex:

tests/utils/test_cache.py

@@ -78,6 +78,29 @@ def test_optimized_query_cache(tmp_path: Path, mocker: MockerFixture):
     assert model._query_renderer._optimized_cache is not None
 
 
+def test_optimized_query_cache_reuses_decoded_entry(tmp_path: Path, mocker: MockerFixture):
+    model = SqlModel(
+        name="test_model",
+        query=parse_one("SELECT a FROM tbl"),
+        mapping_schema={"tbl": {"a": "int"}},
+    )
+    cache = OptimizedQueryCache(tmp_path)
+    assert not cache.with_optimized_query(model)
+
+    get_mock = mocker.spy(cache._file_cache, "get")
+    cache_entry = cache.get_or_create_entry(model)
+    assert cache_entry is not None
+    assert get_mock.call_count == 1
+
+    model._query_renderer._cache = []
+    model._query_renderer._optimized_cache = None
+    cache.with_optimized_query_entry(model, cache_entry)
+
+    # Applying the entry in the parent process must not read and decompress it again.
+    assert get_mock.call_count == 1
+    assert model._query_renderer._optimized_cache is not None
+
+
 def test_optimized_query_cache_missing_rendered_query(tmp_path: Path, mocker: MockerFixture):
     model = SqlModel(
         name="test_model",