Fix: get_snapshots with empty list and dedup (#1573)

Author: tobymao

sqlmesh/core/state_sync/base.py

@@ -323,7 +323,7 @@ def delete_expired_environments(self) -> t.List[Environment]:
 
     @abc.abstractmethod
     def unpause_snapshots(
-        self, snapshots: t.Iterable[SnapshotInfoLike], unpaused_dt: TimeLike
+        self, snapshots: t.Collection[SnapshotInfoLike], unpaused_dt: TimeLike
     ) -> None:
         """Unpauses target snapshots.
 

sqlmesh/core/state_sync/cache.py

@@ -50,7 +50,7 @@ def get_snapshots(
         snapshot_ids: t.Optional[t.Iterable[SnapshotIdLike]],
         hydrate_seeds: bool = False,
     ) -> t.Dict[SnapshotId, Snapshot]:
-        if not snapshot_ids:
+        if snapshot_ids is None:
             return self.state_sync.get_snapshots(snapshot_ids, hydrate_seeds)
 
         existing = {}
@@ -144,7 +144,7 @@ def remove_interval(
         self.state_sync.remove_interval(snapshot_intervals, execution_time, remove_shared_versions)
 
     def unpause_snapshots(
-        self, snapshots: t.Iterable[SnapshotInfoLike], unpaused_dt: TimeLike
+        self, snapshots: t.Collection[SnapshotInfoLike], unpaused_dt: TimeLike
     ) -> None:
         self.snapshot_cache.clear()
         self.state_sync.unpause_snapshots(snapshots, unpaused_dt)

sqlmesh/core/state_sync/common.py

@@ -224,7 +224,7 @@ def _is_snapshot_used(snapshot: Snapshot) -> bool:
 
     @transactional()
     def unpause_snapshots(
-        self, snapshots: t.Iterable[SnapshotInfoLike], unpaused_dt: TimeLike
+        self, snapshots: t.Collection[SnapshotInfoLike], unpaused_dt: TimeLike
     ) -> None:
         current_ts = now()
 
@@ -335,7 +335,7 @@ def _get_snapshots(
     @abc.abstractmethod
     def _get_snapshots_with_same_version(
         self,
-        snapshots: t.Iterable[SnapshotNameVersionLike],
+        snapshots: t.Collection[SnapshotNameVersionLike],
         lock_for_update: bool = False,
     ) -> t.List[Snapshot]:
         """Fetches all snapshots that share the same version as the snapshots.

sqlmesh/core/state_sync/engine_adapter.py

@@ -363,7 +363,11 @@ def _get_snapshots(
         query = (
             exp.select(exp.column("snapshot", table="snapshots"))
             .from_(exp.to_table(self.snapshots_table).as_("snapshots"))
-            .where(self._snapshot_id_filter(snapshot_ids, "snapshots") if snapshot_ids else None)
+            .where(
+                None
+                if snapshot_ids is None
+                else self._snapshot_id_filter(snapshot_ids, "snapshots")
+            )
         )
         if hydrate_seeds:
             query = query.select(exp.column("content", table="seeds")).join(
@@ -411,7 +415,7 @@ def _get_snapshots(
 
     def _get_snapshots_with_same_version(
         self,
-        snapshots: t.Iterable[SnapshotNameVersionLike],
+        snapshots: t.Collection[SnapshotNameVersionLike],
         lock_for_update: bool = False,
     ) -> t.List[Snapshot]:
         """Fetches all snapshots that share the same version as the snapshots.
@@ -906,7 +910,11 @@ def map_data_versions(
     def _snapshot_id_filter(
         self, snapshot_ids: t.Iterable[SnapshotIdLike], alias: t.Optional[str] = None
     ) -> t.Union[exp.In, exp.Boolean, exp.Condition]:
-        if not snapshot_ids:
+        name_identifiers = {
+            (snapshot_id.name, snapshot_id.identifier) for snapshot_id in snapshot_ids
+        }
+
+        if not name_identifiers:
             return exp.false()
         elif self.engine_adapter.SUPPORTS_TUPLE_IN:
             return t.cast(
@@ -917,22 +925,24 @@ def _snapshot_id_filter(
                         exp.column("identifier", table=alias),
                     )
                 ),
-            ).isin(*[(snapshot_id.name, snapshot_id.identifier) for snapshot_id in snapshot_ids])
+            ).isin(*name_identifiers)
         else:
             return exp.or_(
                 *[
                     exp.and_(
-                        exp.column("name", table=alias).eq(snapshot_id.name),
-                        exp.column("identifier", table=alias).eq(snapshot_id.identifier),
+                        exp.column("name", table=alias).eq(name),
+                        exp.column("identifier", table=alias).eq(identifier),
                     )
-                    for snapshot_id in snapshot_ids
+                    for name, identifier in name_identifiers
                 ]
             )
 
     def _snapshot_name_version_filter(
         self, snapshot_name_versions: t.Iterable[SnapshotNameVersionLike], alias: str = "snapshots"
     ) -> t.Union[exp.In, exp.Boolean, exp.Condition]:
-        if not snapshot_name_versions:
+        name_versions = {(s.name, s.version) for s in snapshot_name_versions}
+
+        if not name_versions:
             return exp.false()
         elif self.engine_adapter.SUPPORTS_TUPLE_IN:
             return t.cast(
@@ -943,20 +953,15 @@ def _snapshot_name_version_filter(
                         exp.column("version", table=alias),
                     )
                 ),
-            ).isin(
-                *[
-                    (snapshot_name_version.name, snapshot_name_version.version)
-                    for snapshot_name_version in snapshot_name_versions
-                ]
-            )
+            ).isin(*name_versions)
         else:
             return exp.or_(
                 *[
                     exp.and_(
-                        exp.column("name", table=alias).eq(snapshot_name_version.name),
-                        exp.column("version", table=alias).eq(snapshot_name_version.version),
+                        exp.column("name", table=alias).eq(name),
+                        exp.column("version", table=alias).eq(version),
                     )
-                    for snapshot_name_version in snapshot_name_versions
+                    for name, version in name_versions
                 ]
             )
 

sqlmesh/schedulers/airflow/state_sync.py

@@ -263,7 +263,7 @@ def delete_expired_environments(self) -> t.List[Environment]:
         )
 
     def unpause_snapshots(
-        self, snapshots: t.Iterable[SnapshotInfoLike], unpaused_dt: TimeLike
+        self, snapshots: t.Collection[SnapshotInfoLike], unpaused_dt: TimeLike
     ) -> None:
         """Unpauses target snapshots.
 

tests/core/test_state_sync.py

@@ -1318,3 +1318,13 @@ def test_max_interval_end_for_environment(
     assert state_sync.max_interval_end_for_environment(environment_name) == to_timestamp(
         "2023-01-03"
     )
+
+
+def test_get_snapshots(mocker):
+    mock = mocker.MagicMock()
+    cache = CachingStateSync(mock)
+    cache.get_snapshots([])
+    mock.get_snapshots.assert_not_called()
+
+    cache.get_snapshots(None)
+    mock.get_snapshots.assert_called()