Repository: TobikoData/sqlmesh
Pull request #3041 (2024-08-23 19:57:02)

Fix: Don't cache snapshot intervals
Intervals should never be cached so that they can be hydrated correctly after snapshots are loaded from cache.


Patch:
diff --git a/sqlmesh/core/snapshot/cache.py b/sqlmesh/core/snapshot/cache.py
index b15fdbd3..af573a80 100644
--- a/sqlmesh/core/snapshot/cache.py
+++ b/sqlmesh/core/snapshot/cache.py
@@ -37,6 +37,8 @@ class SnapshotCache:
                 if snapshot.is_model:
                     self._optimized_query_cache.with_optimized_query(snapshot.model)
                 self._update_node_hash_cache(snapshot)
+                snapshot.intervals = []
+                snapshot.dev_intervals = []
                 snapshots[s_id] = snapshot
                 cache_hits.add(s_id)
 
diff --git a/sqlmesh/core/snapshot/definition.py b/sqlmesh/core/snapshot/definition.py
index cbdf352d..1a56d7ca 100644
--- a/sqlmesh/core/snapshot/definition.py
+++ b/sqlmesh/core/snapshot/definition.py
@@ -1207,6 +1207,14 @@ class Snapshot(PydanticModel, SnapshotInfoMixin):
         if not self.version:
             raise SQLMeshError(f"Snapshot {self.snapshot_id} has not been versioned yet.")
 
+    def __getstate__(self) -> t.Dict[t.Any, t.Any]:
+        state = super().__getstate__()
+        state["__dict__"] = state["__dict__"].copy()
+        # Don't store intervals.
+        state["__dict__"]["intervals"] = []
+        state["__dict__"]["dev_intervals"] = []
+        return state
+
 
 class SnapshotTableCleanupTask(PydanticModel):
     snapshot: SnapshotTableInfo


Test patch:
diff --git a/tests/core/test_snapshot.py b/tests/core/test_snapshot.py
index 99bc73fe..332a8148 100644
--- a/tests/core/test_snapshot.py
+++ b/tests/core/test_snapshot.py
@@ -1,3 +1,4 @@
+import pickle
 import json
 import typing as t
 from copy import deepcopy
@@ -1931,6 +1932,8 @@ def test_snapshot_cache(make_snapshot, tmp_path):
             query=parse_one("SELECT 1"),
         )
     )
+    snapshot.add_interval("2024-01-01", "2024-01-02")
+    snapshot.add_interval("2024-01-01", "2024-01-02", is_dev=True)
 
     loader_called_times = 0
 
@@ -1957,6 +1960,8 @@ def test_snapshot_cache(make_snapshot, tmp_path):
     assert cached_snapshot.model._query_renderer._optimized_cache is not None
     assert cached_snapshot.model._data_hash is not None
     assert cached_snapshot.model._metadata_hash is not None
+    assert not cached_snapshot.intervals
+    assert not cached_snapshot.dev_intervals
 
     cache.clear()
     assert cache.get_or_load([snapshot.snapshot_id], _loader) == (
@@ -1964,3 +1969,23 @@ def test_snapshot_cache(make_snapshot, tmp_path):
         set(),
     )
     assert loader_called_times == 2
+
+
+def test_snapshot_pickle_intervals(make_snapshot):
+    snapshot = make_snapshot(
+        SqlModel(
+            name="test_model_name",
+            query=parse_one("SELECT 1"),
+        )
+    )
+    snapshot.add_interval("2023-01-01", "2023-01-02")
+    snapshot.add_interval("2023-01-01", "2023-01-02", is_dev=True)
+
+    assert len(snapshot.intervals) > 0
+    assert len(snapshot.dev_intervals) > 0
+
+    loaded_snapshot = pickle.loads(pickle.dumps(snapshot))
+    assert not loaded_snapshot.intervals
+    assert not loaded_snapshot.dev_intervals
+    assert len(snapshot.intervals) > 0
+    assert len(snapshot.dev_intervals) > 0