Repository: TobikoData/sqlmesh
Pull request #2205 (2024-03-01 23:51:12)

Fix: Set the 'unrestorable' flag even when the snapshot has already been paused
In the case when a BREAKING snapshot has been superseded by an INDIRECT_NON_BREAKING snapshot which in turn has been subsequently replaced by a FORWARD_ONLY snapshot, we end up in a situation when the original snapshot has already been paused (because of the INDIRECT_NON_BREAKING update) but the `unrestorable` flag is not set for it despite the newest update being `FORWARD_ONLY`.


Patch:
diff --git a/sqlmesh/core/state_sync/common.py b/sqlmesh/core/state_sync/common.py
index a440561b..4bd8cf35 100644
--- a/sqlmesh/core/state_sync/common.py
+++ b/sqlmesh/core/state_sync/common.py
@@ -240,11 +240,23 @@ class CommonStateSyncMixin(StateSync):
                         [(snapshot, snapshot.get_removal_interval(effective_from_ts, current_ts))]
                     )
 
+                update_required = False
+
                 if snapshot.unpaused_ts:
                     logger.info("Pausing snapshot %s", snapshot.snapshot_id)
                     snapshot.set_unpaused_ts(None)
-                    if not snapshot.is_forward_only and target_snapshot.is_forward_only:
-                        snapshot.unrestorable = True
+                    update_required = True
+
+                if (
+                    not snapshot.is_forward_only
+                    and target_snapshot.is_forward_only
+                    and not snapshot.unrestorable
+                ):
+                    logger.info("Marking snapshot %s as unrestorable", snapshot.snapshot_id)
+                    snapshot.unrestorable = True
+                    update_required = True
+
+                if update_required:
                     self._update_snapshot(snapshot)
 
     def _ensure_no_gaps(


Test patch:
diff --git a/tests/core/test_state_sync.py b/tests/core/test_state_sync.py
index b687e5ba..2fbff176 100644
--- a/tests/core/test_state_sync.py
+++ b/tests/core/test_state_sync.py
@@ -1173,6 +1173,70 @@ def test_unpause_snapshots(state_sync: EngineAdapterStateSync, make_snapshot: t.
     assert not actual_snapshots[new_snapshot.snapshot_id].unrestorable
 
 
+def test_unrestorable_snapshot(state_sync: EngineAdapterStateSync, make_snapshot: t.Callable):
+    snapshot = make_snapshot(
+        SqlModel(
+            name="test_snapshot",
+            query=parse_one("select 1, ds"),
+            cron="@daily",
+        ),
+    )
+    snapshot.categorize_as(SnapshotChangeCategory.BREAKING)
+    snapshot.version = "a"
+
+    assert not snapshot.unpaused_ts
+    state_sync.push_snapshots([snapshot])
+
+    unpaused_dt = "2022-01-01"
+    state_sync.unpause_snapshots([snapshot], unpaused_dt)
+
+    actual_snapshot = state_sync.get_snapshots([snapshot])[snapshot.snapshot_id]
+    assert actual_snapshot.unpaused_ts
+    assert actual_snapshot.unpaused_ts == to_timestamp(unpaused_dt)
+
+    new_indirect_non_breaking_snapshot = make_snapshot(
+        SqlModel(name="test_snapshot", query=parse_one("select 2, ds"), cron="@daily")
+    )
+    new_indirect_non_breaking_snapshot.categorize_as(SnapshotChangeCategory.INDIRECT_NON_BREAKING)
+    new_indirect_non_breaking_snapshot.version = "a"
+
+    assert not new_indirect_non_breaking_snapshot.unpaused_ts
+    state_sync.push_snapshots([new_indirect_non_breaking_snapshot])
+    state_sync.unpause_snapshots([new_indirect_non_breaking_snapshot], unpaused_dt)
+
+    actual_snapshots = state_sync.get_snapshots([snapshot, new_indirect_non_breaking_snapshot])
+    assert not actual_snapshots[snapshot.snapshot_id].unpaused_ts
+    assert actual_snapshots[
+        new_indirect_non_breaking_snapshot.snapshot_id
+    ].unpaused_ts == to_timestamp(unpaused_dt)
+
+    assert not actual_snapshots[snapshot.snapshot_id].unrestorable
+    assert not actual_snapshots[new_indirect_non_breaking_snapshot.snapshot_id].unrestorable
+
+    new_forward_only_snapshot = make_snapshot(
+        SqlModel(name="test_snapshot", query=parse_one("select 3, ds"), cron="@daily")
+    )
+    new_forward_only_snapshot.categorize_as(SnapshotChangeCategory.FORWARD_ONLY)
+    new_forward_only_snapshot.version = "a"
+
+    assert not new_forward_only_snapshot.unpaused_ts
+    state_sync.push_snapshots([new_forward_only_snapshot])
+    state_sync.unpause_snapshots([new_forward_only_snapshot], unpaused_dt)
+
+    actual_snapshots = state_sync.get_snapshots(
+        [snapshot, new_indirect_non_breaking_snapshot, new_forward_only_snapshot]
+    )
+    assert not actual_snapshots[snapshot.snapshot_id].unpaused_ts
+    assert not actual_snapshots[new_indirect_non_breaking_snapshot.snapshot_id].unpaused_ts
+    assert actual_snapshots[new_forward_only_snapshot.snapshot_id].unpaused_ts == to_timestamp(
+        unpaused_dt
+    )
+
+    assert actual_snapshots[snapshot.snapshot_id].unrestorable
+    assert actual_snapshots[new_indirect_non_breaking_snapshot.snapshot_id].unrestorable
+    assert not actual_snapshots[new_forward_only_snapshot.snapshot_id].unrestorable
+
+
 def test_unpause_snapshots_remove_intervals(
     state_sync: EngineAdapterStateSync, make_snapshot: t.Callable
 ):