Fix: Invalidate the cached model's rendered query after updating the mapping schema (#2884)

Author: izeigerman

sqlmesh/core/model/cache.py

@@ -14,7 +14,7 @@
 
 class SqlModelCacheEntry(PydanticModel):
     model: SqlModel
-    rendered_query: t.Optional[exp.Expression] = None
+    full_depends_on: t.Set[str]
 
 
 class ModelCache:
@@ -46,21 +46,21 @@ def get_or_load(self, name: str, entry_id: str = "", *, loader: t.Callable[[], M
         cache_entry = self._file_cache.get(name, entry_id)
         if cache_entry:
             model = cache_entry.model
-            model._query_renderer.update_cache(cache_entry.rendered_query, optimized=False)
+            model._full_depends_on = cache_entry.full_depends_on
             return model
 
         loaded_model = loader()
         if isinstance(loaded_model, SqlModel):
             new_entry = SqlModelCacheEntry(
-                model=loaded_model, rendered_query=loaded_model.render_query(optimize=False)
+                model=loaded_model, full_depends_on=loaded_model.full_depends_on
             )
             self._file_cache.put(name, entry_id, value=new_entry)
 
         return loaded_model
 
 
 class OptimizedQueryCacheEntry(PydanticModel):
-    optimized_rendered_query: exp.Expression
+    optimized_rendered_query: t.Optional[exp.Expression]
 
 
 class OptimizedQueryCache:
@@ -85,23 +85,29 @@ def with_optimized_query(self, model: Model) -> bool:
         if not isinstance(model, SqlModel):
             return False
 
-        unoptimized_query = model.render_query(optimize=False)
-        if unoptimized_query is None:
-            return False
-
         hash_data = _mapping_schema_hash_data(model.mapping_schema)
-        hash_data.append(gen(unoptimized_query))
+        hash_data.append(gen(model.query))
+        hash_data.append(str([(k, v) for k, v in model.sorted_python_env]))
+        hash_data.extend(model.jinja_macros.data_hash_values)
+
         name = f"{model.name}_{crc32(hash_data)}"
         cache_entry = self._file_cache.get(name)
 
         if cache_entry:
-            model._query_renderer.update_cache(cache_entry.optimized_rendered_query, optimized=True)
+            if cache_entry.optimized_rendered_query:
+                model._query_renderer.update_cache(
+                    cache_entry.optimized_rendered_query, optimized=True
+                )
+            else:
+                # 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.
+                model._query_renderer.update_cache(None, optimized=False)
             return True
 
-        optimized_query = model.render_query(optimize=True)
-        if optimized_query is not None:
-            new_entry = OptimizedQueryCacheEntry(optimized_rendered_query=optimized_query)
-            self._file_cache.put(name, value=new_entry)
+        optimized_query = model.render_query()
+        new_entry = OptimizedQueryCacheEntry(optimized_rendered_query=optimized_query)
+        self._file_cache.put(name, value=new_entry)
 
         return False
 

sqlmesh/core/model/definition.py

@@ -115,6 +115,8 @@ class _Model(ModelMeta, frozen=True):
     jinja_macros: JinjaMacroRegistry = JinjaMacroRegistry()
     mapping_schema: t.Dict[str, t.Any] = {}
 
+    _full_depends_on: t.Optional[t.Set[str]] = None
+
     _expressions_validator = expression_validator
 
     def render(
@@ -560,7 +562,7 @@ def depends_on(self) -> t.Set[str]:
         Returns:
             A list of all the upstream table names.
         """
-        return self._full_depends_on - {self.fqn}
+        return self.full_depends_on - {self.fqn}
 
     @property
     def columns_to_types(self) -> t.Optional[t.Dict[str, exp.DataType]]:
@@ -624,7 +626,7 @@ def is_seed(self) -> bool:
 
     @cached_property
     def depends_on_self(self) -> bool:
-        return self.fqn in self._full_depends_on
+        return self.fqn in self.full_depends_on
 
     @property
     def forward_only(self) -> bool:
@@ -834,17 +836,19 @@ def _additional_metadata(self) -> t.List[str]:
 
         return additional_metadata
 
-    @cached_property
-    def _full_depends_on(self) -> t.Set[str]:
-        depends_on = self.depends_on_ or set()
-
-        query = self.render_query(optimize=False)
-        if query is not None:
-            depends_on |= d.find_tables(
-                query, default_catalog=self.default_catalog, dialect=self.dialect
-            )
+    @property
+    def full_depends_on(self) -> t.Set[str]:
+        if not self._full_depends_on:
+            depends_on = self.depends_on_ or set()
+
+            query = self.render_query(optimize=False)
+            if query is not None:
+                depends_on |= d.find_tables(
+                    query, default_catalog=self.default_catalog, dialect=self.dialect
+                )
+            self._full_depends_on = depends_on
 
-        return depends_on
+        return self._full_depends_on
 
 
 class _SqlBasedModel(_Model):
@@ -1092,7 +1096,7 @@ def column_descriptions(self) -> t.Dict[str, str]:
         if self.column_descriptions_ is not None:
             return self.column_descriptions_
 
-        query = self.render_query(optimize=False)
+        query = self.render_query()
         if query is None:
             return {}
 

sqlmesh/utils/cache.py

@@ -119,7 +119,7 @@ def put(self, name: str, entry_id: str = "", *, value: T) -> None:
             raise SQLMeshError(f"Cache path '{self._path}' is not a directory.")
 
         with gzip.open(self._cache_entry_path(name, entry_id), "wb", compresslevel=1) as fd:
-            pickle.dump(value.dict(), fd)
+            pickle.dump(value.dict(exclude_none=False), fd)
 
     def _cache_entry_path(self, name: str, entry_id: str = "") -> Path:
         entry_file_name = "__".join(p for p in (self._cache_version, name, entry_id) if p)

tests/core/test_integration.py

@@ -6,15 +6,23 @@
 
 import pandas as pd
 import pytest
+from pathlib import Path
 from freezegun import freeze_time
 from pytest_mock.plugin import MockerFixture
 from sqlglot import exp
 from sqlglot.expressions import DataType
 
 from sqlmesh import CustomMaterialization
+from sqlmesh.cli.example_project import init_example_project
 from sqlmesh.core import constants as c
 from sqlmesh.core import dialect as d
-from sqlmesh.core.config import AutoCategorizationMode
+from sqlmesh.core.config import (
+    AutoCategorizationMode,
+    Config,
+    GatewayConfig,
+    ModelDefaultsConfig,
+    DuckDBConnectionConfig,
+)
 from sqlmesh.core.console import Console
 from sqlmesh.core.context import Context
 from sqlmesh.core.engine_adapter import EngineAdapter
@@ -1391,6 +1399,43 @@ def test_restatement_plan_ignores_changes(init_and_plan_context: t.Callable):
     context.apply(plan)
 
 
+def test_plan_twice_with_star_macro_yields_no_diff(tmp_path: Path):
+    init_example_project(tmp_path, dialect="duckdb")
+
+    star_model_definition = """
+        MODEL (
+          name sqlmesh_example.star_model,
+          kind FULL
+        );
+
+        SELECT @STAR(sqlmesh_example.full_model) FROM sqlmesh_example.full_model
+    """
+
+    star_model_path = tmp_path / "models" / "star_model.sql"
+    star_model_path.write_text(star_model_definition)
+
+    db_path = str(tmp_path / "db.db")
+    config = Config(
+        gateways={"main": GatewayConfig(connection=DuckDBConnectionConfig(database=db_path))},
+        model_defaults=ModelDefaultsConfig(dialect="duckdb"),
+    )
+    context = Context(paths=tmp_path, config=config)
+    context.plan(auto_apply=True, no_prompts=True)
+
+    # Instantiate new context to remove caches etc
+    new_context = Context(paths=tmp_path, config=config)
+
+    star_model = new_context.get_model("sqlmesh_example.star_model")
+    assert (
+        star_model.render_query_or_raise().sql()
+        == 'SELECT CAST("full_model"."item_id" AS INT) AS "item_id", CAST("full_model"."num_orders" AS BIGINT) AS "num_orders" FROM "db"."sqlmesh_example"."full_model" AS "full_model"'
+    )
+
+    new_plan = new_context.plan(no_prompts=True)
+    assert not new_plan.has_changes
+    assert not new_plan.new_snapshots
+
+
 @pytest.mark.parametrize(
     "context_fixture",
     ["sushi_context", "sushi_dbt_context", "sushi_test_dbt_context", "sushi_no_default_catalog"],

tests/utils/test_cache.py

@@ -49,4 +49,33 @@ def test_optimized_query_cache(tmp_path: Path, mocker: MockerFixture):
     cache = OptimizedQueryCache(tmp_path)
 
     assert not cache.with_optimized_query(model)
+
+    model._query_renderer._cache = []
+    model._query_renderer._optimized_cache = None
+
+    assert cache.with_optimized_query(model)
+
+    assert not model._query_renderer._cache
+    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",
+        query=parse_one("SELECT a FROM tbl"),
+        mapping_schema={"tbl": {"a": "int"}},
+    )
+    render_mock = mocker.patch.object(model._query_renderer, "render")
+    render_mock.return_value = None
+
+    cache = OptimizedQueryCache(tmp_path)
+
+    assert not cache.with_optimized_query(model)
+
+    model._query_renderer._cache = []
+    model._query_renderer._optimized_cache = None
+
     assert cache.with_optimized_query(model)
+
+    assert model._query_renderer._cache == [None]
+    assert model._query_renderer._optimized_cache is None