Fix: Plan application should succeeed even when the snapshot query in state is unrenderable (#4235)

Author: izeigerman

sqlmesh/core/context.py

@@ -2144,6 +2144,8 @@ def table_name(self, model_name: str, dev: bool) -> str:
     def clear_caches(self) -> None:
         for path in self.configs:
             rmtree(path / c.CACHE)
+        if isinstance(self.state_sync, CachingStateSync):
+            self.state_sync.clear_cache()
 
     def export_state(
         self,

sqlmesh/core/model/definition.py

@@ -1337,7 +1337,11 @@ def columns_to_types(self) -> t.Optional[t.Dict[str, exp.DataType]]:
         if self.columns_to_types_ is not None:
             self._columns_to_types = self.columns_to_types_
         elif self._columns_to_types is None:
-            query = self._query_renderer.render()
+            try:
+                query = self._query_renderer.render()
+            except Exception:
+                logger.exception("Failed to render query for model %s", self.fqn)
+                return None
 
             if query is None:
                 return None

sqlmesh/core/snapshot/cache.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import logging
 import typing as t
 
 from pathlib import Path
@@ -13,6 +14,9 @@
 from sqlmesh.utils.cache import FileCache
 
 
+logger = logging.getLogger(__name__)
+
+
 class SnapshotCache:
     def __init__(self, path: Path):
         self._snapshot_cache: FileCache[Snapshot] = FileCache(path, prefix="snapshot")
@@ -70,7 +74,12 @@ def get_or_load(
             self._update_node_hash_cache(snapshot)
 
             if snapshot.is_model and c.MAX_FORK_WORKERS == 1:
-                self._optimized_query_cache.with_optimized_query(snapshot.model)
+                try:
+                    self._optimized_query_cache.with_optimized_query(snapshot.model)
+                except Exception:
+                    logger.exception(
+                        "Failed to cache optimized query for snapshot %s", snapshot.snapshot_id
+                    )
 
             self.put(snapshot)
 
@@ -82,10 +91,13 @@ def put(self, snapshot: Snapshot) -> None:
         if self._snapshot_cache.exists(entry_name):
             return
 
-        if snapshot.is_model:
-            # make sure we preload full_depends_on
-            snapshot.model.full_depends_on
-        self._snapshot_cache.put(entry_name, value=snapshot)
+        try:
+            if snapshot.is_model:
+                # make sure we preload full_depends_on
+                snapshot.model.full_depends_on
+            self._snapshot_cache.put(entry_name, value=snapshot)
+        except Exception:
+            logger.exception("Failed to cache snapshot %s", snapshot.snapshot_id)
 
     def clear(self) -> None:
         self._snapshot_cache.clear()

sqlmesh/core/state_sync/cache.py

@@ -143,3 +143,6 @@ def unpause_snapshots(
     ) -> None:
         self.snapshot_cache.clear()
         self.state_sync.unpause_snapshots(snapshots, unpaused_dt)
+
+    def clear_cache(self) -> None:
+        self.snapshot_cache.clear()

tests/core/test_integration.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import typing as t
+import json
 from collections import Counter
 from datetime import timedelta
 from unittest import mock
@@ -3995,6 +3996,59 @@ def test_empty_backfill_new_model(init_and_plan_context: t.Callable):
                 assert snapshot.intervals[-1][1] <= to_timestamp("2023-01-08")
 
 
+@time_machine.travel("2023-01-08 15:00:00 UTC")
+@pytest.mark.parametrize("forward_only", [False, True])
+def test_plan_repairs_unrenderable_snapshot_state(
+    init_and_plan_context: t.Callable, forward_only: bool
+):
+    context, plan = init_and_plan_context("examples/sushi")
+    context.apply(plan)
+
+    target_snapshot = context.get_snapshot("sushi.waiter_revenue_by_day")
+    assert target_snapshot
+
+    # Manually corrupt the snapshot's query
+    raw_snapshot = context.state_sync.state_sync.engine_adapter.fetchone(
+        f"SELECT snapshot FROM sqlmesh._snapshots WHERE name = '{target_snapshot.name}' AND identifier = '{target_snapshot.identifier}'"
+    )[0]  # type: ignore
+    parsed_snapshot = json.loads(raw_snapshot)
+    parsed_snapshot["node"]["query"] = "SELECT @missing_macro()"
+    context.state_sync.state_sync.engine_adapter.update_table(
+        "sqlmesh._snapshots",
+        {"snapshot": json.dumps(parsed_snapshot)},
+        f"name = '{target_snapshot.name}' AND identifier = '{target_snapshot.identifier}'",
+    )
+
+    context.clear_caches()
+
+    target_snapshot_in_state = context.state_sync.get_snapshots([target_snapshot.snapshot_id])[
+        target_snapshot.snapshot_id
+    ]
+    with pytest.raises(Exception):
+        target_snapshot_in_state.model.render_query_or_raise()
+
+    # Repair the snapshot by creating a new version of it
+    context.upsert_model(target_snapshot.model.name, stamp="repair")
+    target_snapshot = context.get_snapshot(target_snapshot.name)
+
+    plan_builder = context.plan_builder("prod", forward_only=forward_only)
+    plan = plan_builder.build()
+    assert plan.directly_modified == {target_snapshot.snapshot_id}
+    if not forward_only:
+        assert {i.snapshot_id for i in plan.missing_intervals} == {target_snapshot.snapshot_id}
+        plan_builder.set_choice(target_snapshot, SnapshotChangeCategory.NON_BREAKING)
+        plan = plan_builder.build()
+
+    context.apply(plan)
+
+    context.clear_caches()
+    assert context.get_snapshot(target_snapshot.name).model.render_query_or_raise()
+    target_snapshot_in_state = context.state_sync.get_snapshots([target_snapshot.snapshot_id])[
+        target_snapshot.snapshot_id
+    ]
+    assert target_snapshot_in_state.model.render_query_or_raise()
+
+
 @time_machine.travel("2023-01-08 15:00:00 UTC")
 def test_dbt_requirements(sushi_dbt_context: Context):
     assert set(sushi_dbt_context.requirements) == {"dbt-core", "dbt-duckdb"}