PR feedback

Author: erindru

sqlmesh/core/plan/common.py

@@ -134,12 +134,20 @@ def identify_restatement_intervals_across_snapshot_versions(
     unique_snapshot_names = set(snapshot_id.name for snapshot_id in snapshot_intervals_to_clear)
 
     current_ts = current_ts or now_timestamp()
-    all_matching_snapshot_ids = state_reader.get_snapshot_ids_by_names(
-        snapshot_names=unique_snapshot_names, current_ts=current_ts, exclude_expired=True
-    )
+    all_matching_non_prod_snapshots = {
+        s.snapshot_id: s
+        for s in state_reader.get_snapshots_by_names(
+            snapshot_names=unique_snapshot_names, current_ts=current_ts, exclude_expired=True
+        )
+        # Don't clear intervals for a snapshot if it shares the same physical version with prod.
+        # Otherwise, prod will be affected by what should be a dev operation
+        if s.name_version not in prod_name_versions
+    }
 
     # identify the ones that we havent picked up yet, which are the ones that dont exist in any environment
-    if remaining_snapshot_ids := all_matching_snapshot_ids.difference(snapshot_intervals_to_clear):
+    if remaining_snapshot_ids := set(all_matching_non_prod_snapshots).difference(
+        snapshot_intervals_to_clear
+    ):
         # these snapshot id's exist in isolation and may be related to a downstream dependency of the :prod_restatements,
         # rather than directly related, so we can't simply look up the interval to clear based on :prod_restatements.
         # To figure out the interval that should be cleared, we can match to the existing list based on name
@@ -156,18 +164,19 @@ def identify_restatement_intervals_across_snapshot_versions(
 
             snapshot_name_to_widest_interval[s_id.name] = (next_start, next_end)
 
+        # we need to fetch full Snapshot's to get access to the SnapshotTableInfo objects
+        # required by StateSync.remove_intervals()
+        # but at this point we have minimized the list by excluding the ones that are already present in prod
+        # and also excluding the ones we have already matched earlier while traversing the environment DAGs
         remaining_snapshots = state_reader.get_snapshots(snapshot_ids=remaining_snapshot_ids)
         for remaining_snapshot_id, remaining_snapshot in remaining_snapshots.items():
-            # Don't clear intervals for a snapshot if it shares the same physical version with prod.
-            # Otherwise, prod will be affected by what should be a dev operation
-            if remaining_snapshot.name_version in prod_name_versions:
-                continue
-
             snapshot_intervals_to_clear[remaining_snapshot_id] = SnapshotIntervalClearRequest(
                 table_info=remaining_snapshot.table_info,
                 interval=snapshot_name_to_widest_interval[remaining_snapshot_id.name],
             )
 
+        loaded_snapshots.update(remaining_snapshots)
+
     # for any affected full_history_restatement_only snapshots, we need to widen the intervals being restated to
     # include the whole time range for that snapshot. This requires a call to state to load the full snapshot record,
     # so we only do it if necessary