diff --git a/sqlmesh/core/model/cache.py b/sqlmesh/core/model/cache.py index 1f038c5d79..cba1e05c00 100644 --- a/sqlmesh/core/model/cache.py +++ b/sqlmesh/core/model/cache.py @@ -112,15 +112,7 @@ 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) @@ -128,6 +120,42 @@ def with_optimized_query(self, model: Model, name: t.Optional[str] = None) -> bo 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,7 +168,7 @@ 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( @@ -148,6 +176,7 @@ def _put(self, name: str, model: SqlModel) -> None: 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 +226,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 +242,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, diff --git a/sqlmesh/core/model/schema.py b/sqlmesh/core/model/schema.py index e29cacade0..0d55859fbf 100644 --- a/sqlmesh/core/model/schema.py +++ b/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: diff --git a/tests/utils/test_cache.py b/tests/utils/test_cache.py index e6e041e30a..b151981b72 100644 --- a/tests/utils/test_cache.py +++ b/tests/utils/test_cache.py @@ -62,6 +62,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",