Feat(state_sync): Add the ability to fetch all versions of a snapshot by name

Author: erindru

sqlmesh/core/state_sync/base.py

@@ -97,6 +97,24 @@ def get_snapshots(
             A dictionary of snapshot ids to snapshots for ones that could be found.
         """
 
+    @abc.abstractmethod
+    def get_snapshot_ids_by_names(
+        self,
+        snapshot_names: t.Iterable[str],
+        current_ts: t.Optional[int] = None,
+        exclude_expired: bool = True,
+    ) -> t.Set[SnapshotId]:
+        """Return the snapshot id's for all versions of the specified snapshot names.
+
+        Args:
+            snapshot_names: Iterable of snapshot names to fetch all snapshot id's for
+            current_ts: Sets the current time for identifying which snapshots have expired so they can be excluded (only relevant if :exclude_expired=True)
+            exclude_expired: Whether or not to return the snapshot id's of expired snapshots in the result
+
+        Returns:
+            A dictionary mapping snapshot names to a list of relevant snapshot id's
+        """
+
     @abc.abstractmethod
     def snapshots_exist(self, snapshot_ids: t.Iterable[SnapshotIdLike]) -> t.Set[SnapshotId]:
         """Checks if multiple snapshots exist in the state sync.

sqlmesh/core/state_sync/db/facade.py

@@ -366,6 +366,16 @@ def get_snapshots(
         Snapshot.hydrate_with_intervals_by_version(snapshots.values(), intervals)
         return snapshots
 
+    def get_snapshot_ids_by_names(
+        self,
+        snapshot_names: t.Iterable[str],
+        current_ts: t.Optional[int] = None,
+        exclude_expired: bool = True,
+    ) -> t.Set[SnapshotId]:
+        return self.snapshot_state.get_snapshot_ids_by_names(
+            snapshot_names=snapshot_names, current_ts=current_ts, exclude_expired=exclude_expired
+        )
+
     @transactional()
     def add_interval(
         self,

sqlmesh/core/state_sync/db/snapshot.py

@@ -308,6 +308,46 @@ def get_snapshots(
         """
         return self._get_snapshots(snapshot_ids)
 
+    def get_snapshot_ids_by_names(
+        self,
+        snapshot_names: t.Iterable[str],
+        current_ts: t.Optional[int] = None,
+        exclude_expired: bool = True,
+    ) -> t.Set[SnapshotId]:
+        """Return the snapshot id's for all versions of the specified snapshot names.
+
+        Args:
+            snapshot_names: Iterable of snapshot names to fetch all snapshot id's for
+            current_ts: Sets the current time for identifying which snapshots have expired so they can be excluded (only relevant if :exclude_expired=True)
+            exclude_expired: Whether or not to return the snapshot id's of expired snapshots in the result
+
+        Returns:
+            A dictionary mapping snapshot names to a list of relevant snapshot id's
+        """
+        if not snapshot_names:
+            return set()
+
+        if exclude_expired:
+            current_ts = current_ts or now_timestamp()
+            unexpired_expr = (exp.column("updated_ts") + exp.column("ttl_ms")) > current_ts
+        else:
+            unexpired_expr = None
+
+        return {
+            SnapshotId(name=name, identifier=identifier)
+            for where in snapshot_name_filter(
+                snapshot_names=snapshot_names,
+                batch_size=self.SNAPSHOT_BATCH_SIZE,
+            )
+            for name, identifier in fetchall(
+                self.engine_adapter,
+                exp.select("name", "identifier")
+                .from_(self.snapshots_table)
+                .where(where)
+                .and_(unexpired_expr),
+            )
+        }
+
     def snapshots_exist(self, snapshot_ids: t.Iterable[SnapshotIdLike]) -> t.Set[SnapshotId]:
         """Checks if snapshots exist.
 

tests/core/state_sync/test_state_sync.py

@@ -3569,3 +3569,81 @@ def test_update_environment_statements(state_sync: EngineAdapterStateSync):
         "@grant_schema_usage()",
         "@grant_select_privileges()",
     ]
+
+
+def test_get_snapshot_ids_by_names(
+    state_sync: EngineAdapterStateSync, make_snapshot: t.Callable[..., Snapshot]
+):
+    assert state_sync.get_snapshot_ids_by_names(snapshot_names=[]) == set()
+
+    snap_a_v1, snap_a_v2 = (
+        make_snapshot(
+            SqlModel(
+                name="a",
+                query=parse_one(f"select {i}, ds"),
+            ),
+            version="a",
+        )
+        for i in range(2)
+    )
+
+    snap_b = make_snapshot(
+        SqlModel(
+            name="b",
+            query=parse_one(f"select 'b' as b, ds"),
+        ),
+        version="b",
+    )
+
+    state_sync.push_snapshots([snap_a_v1, snap_a_v2, snap_b])
+
+    assert state_sync.get_snapshot_ids_by_names(snapshot_names=['"a"']) == {
+        snap_a_v1.snapshot_id,
+        snap_a_v2.snapshot_id,
+    }
+    assert state_sync.get_snapshot_ids_by_names(snapshot_names=['"a"', '"b"']) == {
+        snap_a_v1.snapshot_id,
+        snap_a_v2.snapshot_id,
+        snap_b.snapshot_id,
+    }
+
+
+def test_get_snapshot_ids_by_names_include_expired(
+    state_sync: EngineAdapterStateSync, make_snapshot: t.Callable[..., Snapshot]
+):
+    now_ts = now_timestamp()
+
+    normal_a = make_snapshot(
+        SqlModel(
+            name="a",
+            query=parse_one(f"select 1, ds"),
+        ),
+        version="a",
+    )
+
+    expired_a = make_snapshot(
+        SqlModel(
+            name="a",
+            query=parse_one(f"select 2, ds"),
+        ),
+        version="a",
+        ttl="in 10 seconds",
+    )
+    expired_a.updated_ts = now_ts - (
+        1000 * 15
+    )  # last updated 15 seconds ago, expired 10 seconds from last updated = expired 5 seconds ago
+
+    state_sync.push_snapshots([normal_a, expired_a])
+
+    assert state_sync.get_snapshot_ids_by_names(snapshot_names=['"a"'], current_ts=now_ts) == {
+        normal_a.snapshot_id
+    }
+    assert state_sync.get_snapshot_ids_by_names(snapshot_names=['"a"'], exclude_expired=False) == {
+        normal_a.snapshot_id,
+        expired_a.snapshot_id,
+    }
+
+    # wind back time to 10 seconds ago (before the expired snapshot is expired - it expired 5 seconds ago) to test it stil shows in a normal query
+    assert state_sync.get_snapshot_ids_by_names(
+        snapshot_names=['"a"'], current_ts=(now_ts - (10 * 1000))
+    ) == {normal_a.snapshot_id, expired_a.snapshot_id}