PR feedback

Author: erindru

sqlmesh/core/plan/common.py

@@ -4,7 +4,7 @@
 from dataclasses import dataclass, field
 
 from sqlmesh.core.state_sync import StateReader
-from sqlmesh.core.snapshot import Snapshot, SnapshotId, SnapshotTableInfo, SnapshotNameVersion
+from sqlmesh.core.snapshot import Snapshot, SnapshotId, SnapshotIdAndVersion, SnapshotNameVersion
 from sqlmesh.core.snapshot.definition import Interval
 from sqlmesh.utils.dag import DAG
 from sqlmesh.utils.date import now_timestamp
@@ -41,7 +41,7 @@ def is_breaking_kind_change(old: Snapshot, new: Snapshot) -> bool:
 @dataclass
 class SnapshotIntervalClearRequest:
     # affected snapshot
-    table_info: SnapshotTableInfo
+    snapshot: SnapshotIdAndVersion
 
     # which interval to clear
     interval: Interval
@@ -53,7 +53,7 @@ class SnapshotIntervalClearRequest:
 
     @property
     def snapshot_id(self) -> SnapshotId:
-        return self.table_info.snapshot_id
+        return self.snapshot.snapshot_id
 
     @property
     def sorted_environment_names(self) -> t.List[str]:
@@ -122,7 +122,7 @@ def identify_restatement_intervals_across_snapshot_versions(
                 clear_request = snapshot_intervals_to_clear.get(affected_snapshot.snapshot_id)
                 if not clear_request:
                     clear_request = SnapshotIntervalClearRequest(
-                        table_info=affected_snapshot, interval=interval
+                        snapshot=affected_snapshot.id_and_version, interval=interval
                     )
                     snapshot_intervals_to_clear[affected_snapshot.snapshot_id] = clear_request
 
@@ -164,19 +164,13 @@ 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():
+        for remaining_snapshot_id in remaining_snapshot_ids:
+            remaining_snapshot = all_matching_non_prod_snapshots[remaining_snapshot_id]
             snapshot_intervals_to_clear[remaining_snapshot_id] = SnapshotIntervalClearRequest(
-                table_info=remaining_snapshot.table_info,
+                snapshot=remaining_snapshot,
                 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
@@ -187,7 +181,7 @@ def identify_restatement_intervals_across_snapshot_versions(
         # So for now, these are not considered
         s_id
         for s_id, s in snapshot_intervals_to_clear.items()
-        if s.table_info.full_history_restatement_only
+        if s.snapshot.kind_name and s.snapshot.kind_name.full_history_restatement_only
     ]
     if full_history_restatement_snapshot_ids:
         # only load full snapshot records that we havent already loaded

sqlmesh/core/plan/evaluator.py

@@ -284,7 +284,9 @@ def visit_audit_only_run_stage(
     def visit_restatement_stage(
         self, stage: stages.RestatementStage, plan: EvaluatablePlan
     ) -> None:
-        snapshot_intervals_to_restate = {(s, i) for s, i in stage.snapshot_intervals.items()}
+        snapshot_intervals_to_restate = {
+            (s.id_and_version, i) for s, i in stage.snapshot_intervals.items()
+        }
 
         # Restating intervals on prod plans should mean that the intervals are cleared across
         # all environments, not just the version currently in prod
@@ -294,7 +296,7 @@ def visit_restatement_stage(
         # Without this rule, its possible that promoting a dev table to prod will introduce old data to prod
         snapshot_intervals_to_restate.update(
             {
-                (s.table_info, s.interval)
+                (s.snapshot, s.interval)
                 for s in identify_restatement_intervals_across_snapshot_versions(
                     state_reader=self.state_sync,
                     prod_restatements=plan.restatements,

sqlmesh/core/snapshot/__init__.py

@@ -11,6 +11,7 @@
     SnapshotId as SnapshotId,
     SnapshotIdBatch as SnapshotIdBatch,
     SnapshotIdLike as SnapshotIdLike,
+    SnapshotIdAndVersionLike as SnapshotIdAndVersionLike,
     SnapshotInfoLike as SnapshotInfoLike,
     SnapshotIntervals as SnapshotIntervals,
     SnapshotNameVersion as SnapshotNameVersion,

sqlmesh/core/snapshot/definition.py

@@ -587,6 +587,17 @@ def name_version(self) -> SnapshotNameVersion:
         """Returns the name and version of the snapshot."""
         return SnapshotNameVersion(name=self.name, version=self.version)
 
+    @property
+    def id_and_version(self) -> SnapshotIdAndVersion:
+        return SnapshotIdAndVersion(
+            name=self.name,
+            kind_name=self.kind_name,
+            identifier=self.identifier,
+            version=self.version,
+            dev_version=self.dev_version,
+            fingerprint=self.fingerprint,
+        )
+
 
 class SnapshotIdAndVersion(PydanticModel):
     """A stripped down version of a snapshot that is used in situations where we want to fetch the main fields of the snapshots table
@@ -595,6 +606,7 @@ class SnapshotIdAndVersion(PydanticModel):
 
     name: str
     version: str
+    kind_name: t.Optional[ModelKindName] = None
     dev_version_: t.Optional[str] = Field(alias="dev_version")
     identifier: str
     fingerprint_: t.Union[str, SnapshotFingerprint] = Field(alias="fingerprint")
@@ -603,6 +615,10 @@ class SnapshotIdAndVersion(PydanticModel):
     def snapshot_id(self) -> SnapshotId:
         return SnapshotId(name=self.name, identifier=self.identifier)
 
+    @property
+    def id_and_version(self) -> SnapshotIdAndVersion:
+        return self
+
     @property
     def name_version(self) -> SnapshotNameVersion:
         return SnapshotNameVersion(name=self.name, version=self.version)
@@ -1424,6 +1440,10 @@ def name_version(self) -> SnapshotNameVersion:
         """Returns the name and version of the snapshot."""
         return SnapshotNameVersion(name=self.name, version=self.version)
 
+    @property
+    def id_and_version(self) -> SnapshotIdAndVersion:
+        return self.table_info.id_and_version
+
     @property
     def disable_restatement(self) -> bool:
         """Is restatement disabled for the node"""
@@ -1494,7 +1514,8 @@ class SnapshotTableCleanupTask(PydanticModel):
     dev_table_only: bool
 
 
-SnapshotIdLike = t.Union[SnapshotId, SnapshotTableInfo, SnapshotIdAndVersion, Snapshot]
+SnapshotIdLike = t.Union[SnapshotId, SnapshotIdAndVersion, SnapshotTableInfo, Snapshot]
+SnapshotIdAndVersionLike = t.Union[SnapshotIdAndVersion, SnapshotTableInfo, Snapshot]
 SnapshotInfoLike = t.Union[SnapshotTableInfo, Snapshot]
 SnapshotNameVersionLike = t.Union[
     SnapshotNameVersion, SnapshotTableInfo, SnapshotIdAndVersion, Snapshot

sqlmesh/core/state_sync/base.py

@@ -19,6 +19,7 @@
     Snapshot,
     SnapshotId,
     SnapshotIdLike,
+    SnapshotIdAndVersionLike,
     SnapshotInfoLike,
     SnapshotTableCleanupTask,
     SnapshotTableInfo,
@@ -390,7 +391,7 @@ def remove_state(self, including_backup: bool = False) -> None:
     @abc.abstractmethod
     def remove_intervals(
         self,
-        snapshot_intervals: t.Sequence[t.Tuple[SnapshotInfoLike, Interval]],
+        snapshot_intervals: t.Sequence[t.Tuple[SnapshotIdAndVersionLike, Interval]],
         remove_shared_versions: bool = False,
     ) -> None:
         """Remove an interval from a list of snapshots and sync it to the store.

sqlmesh/core/state_sync/cache.py

@@ -7,6 +7,7 @@
     Snapshot,
     SnapshotId,
     SnapshotIdLike,
+    SnapshotIdAndVersionLike,
     SnapshotInfoLike,
 )
 from sqlmesh.core.snapshot.definition import Interval, SnapshotIntervals
@@ -128,7 +129,7 @@ def add_snapshots_intervals(self, snapshots_intervals: t.Sequence[SnapshotInterv
 
     def remove_intervals(
         self,
-        snapshot_intervals: t.Sequence[t.Tuple[SnapshotInfoLike, Interval]],
+        snapshot_intervals: t.Sequence[t.Tuple[SnapshotIdAndVersionLike, Interval]],
         remove_shared_versions: bool = False,
     ) -> None:
         for s, _ in snapshot_intervals:

sqlmesh/core/state_sync/db/facade.py

@@ -31,6 +31,7 @@
     SnapshotIdAndVersion,
     SnapshotId,
     SnapshotIdLike,
+    SnapshotIdAndVersionLike,
     SnapshotInfoLike,
     SnapshotIntervals,
     SnapshotNameVersion,
@@ -407,7 +408,7 @@ def add_snapshots_intervals(self, snapshots_intervals: t.Sequence[SnapshotInterv
     @transactional()
     def remove_intervals(
         self,
-        snapshot_intervals: t.Sequence[t.Tuple[SnapshotInfoLike, Interval]],
+        snapshot_intervals: t.Sequence[t.Tuple[SnapshotIdAndVersionLike, Interval]],
         remove_shared_versions: bool = False,
     ) -> None:
         self.interval_state.remove_intervals(snapshot_intervals, remove_shared_versions)

sqlmesh/core/state_sync/db/interval.py

@@ -15,10 +15,10 @@
 from sqlmesh.core.snapshot import (
     SnapshotIntervals,
     SnapshotIdLike,
+    SnapshotIdAndVersionLike,
     SnapshotNameVersionLike,
     SnapshotTableCleanupTask,
     SnapshotNameVersion,
-    SnapshotInfoLike,
     Snapshot,
 )
 from sqlmesh.core.snapshot.definition import Interval
@@ -68,11 +68,11 @@ def add_snapshots_intervals(self, snapshots_intervals: t.Sequence[SnapshotInterv
 
     def remove_intervals(
         self,
-        snapshot_intervals: t.Sequence[t.Tuple[SnapshotInfoLike, Interval]],
+        snapshot_intervals: t.Sequence[t.Tuple[SnapshotIdAndVersionLike, Interval]],
         remove_shared_versions: bool = False,
     ) -> None:
         intervals_to_remove: t.Sequence[
-            t.Tuple[t.Union[SnapshotInfoLike, SnapshotIntervals], Interval]
+            t.Tuple[t.Union[SnapshotIdAndVersionLike, SnapshotIntervals], Interval]
         ] = snapshot_intervals
         if remove_shared_versions:
             name_version_mapping = {s.name_version: interval for s, interval in snapshot_intervals}
@@ -431,7 +431,9 @@ def _delete_intervals_by_version(self, targets: t.List[SnapshotTableCleanupTask]
 
 
 def _intervals_to_df(
-    snapshot_intervals: t.Sequence[t.Tuple[t.Union[SnapshotInfoLike, SnapshotIntervals], Interval]],
+    snapshot_intervals: t.Sequence[
+        t.Tuple[t.Union[SnapshotIdAndVersionLike, SnapshotIntervals], Interval]
+    ],
     is_dev: bool,
     is_removed: bool,
 ) -> pd.DataFrame:
@@ -451,7 +453,7 @@ def _intervals_to_df(
 
 
 def _interval_to_df(
-    snapshot: t.Union[SnapshotInfoLike, SnapshotIntervals],
+    snapshot: t.Union[SnapshotIdAndVersionLike, SnapshotIntervals],
     start_ts: int,
     end_ts: int,
     is_dev: bool = False,

sqlmesh/core/state_sync/db/snapshot.py

@@ -337,16 +337,19 @@ def get_snapshots_by_names(
                 name=name,
                 identifier=identifier,
                 version=version,
+                kind_name=kind_name,
                 dev_version=dev_version,
                 fingerprint=fingerprint,
             )
             for where in snapshot_name_filter(
                 snapshot_names=snapshot_names,
                 batch_size=self.SNAPSHOT_BATCH_SIZE,
             )
-            for name, identifier, version, dev_version, fingerprint in fetchall(
+            for name, identifier, version, kind_name, dev_version, fingerprint in fetchall(
                 self.engine_adapter,
-                exp.select("name", "identifier", "version", "dev_version", "fingerprint")
+                exp.select(
+                    "name", "identifier", "version", "kind_name", "dev_version", "fingerprint"
+                )
                 .from_(self.snapshots_table)
                 .where(where)
                 .and_(unexpired_expr),
@@ -661,6 +664,7 @@ def _get_snapshots_with_same_version(
                     "name",
                     "identifier",
                     "version",
+                    "kind_name",
                     "dev_version",
                     "fingerprint",
                 )
@@ -677,10 +681,11 @@ def _get_snapshots_with_same_version(
                 name=name,
                 identifier=identifier,
                 version=version,
+                kind_name=kind_name,
                 dev_version=dev_version,
                 fingerprint=SnapshotFingerprint.parse_raw(fingerprint),
             )
-            for name, identifier, version, dev_version, fingerprint in snapshot_rows
+            for name, identifier, version, kind_name, dev_version, fingerprint in snapshot_rows
         ]
 
 

tests/core/test_snapshot.py

@@ -3567,3 +3567,28 @@ def test_snapshot_id_and_version_fingerprint_lazy_init():
 
     assert isinstance(snapshot.fingerprint_, SnapshotFingerprint)
     assert snapshot.fingerprint == fingerprint
+
+
+def test_snapshot_id_and_version_optional_kind_name():
+    snapshot = SnapshotIdAndVersion(
+        name="a",
+        identifier="1234",
+        version="2345",
+        dev_version=None,
+        fingerprint="",
+    )
+
+    assert snapshot.kind_name is None
+
+    snapshot = SnapshotIdAndVersion(
+        name="a",
+        identifier="1234",
+        version="2345",
+        kind_name="INCREMENTAL_UNMANAGED",
+        dev_version=None,
+        fingerprint="",
+    )
+
+    assert snapshot.kind_name
+    assert snapshot.kind_name.is_incremental_unmanaged
+    assert snapshot.kind_name.full_history_restatement_only