PR feedback

Author: erindru

sqlmesh/core/snapshot/__init__.py

@@ -4,6 +4,7 @@
     Node as Node,
     QualifiedViewName as QualifiedViewName,
     Snapshot as Snapshot,
+    MinimalSnapshot as MinimalSnapshot,
     SnapshotChangeCategory as SnapshotChangeCategory,
     SnapshotDataVersion as SnapshotDataVersion,
     SnapshotFingerprint as SnapshotFingerprint,

sqlmesh/core/snapshot/definition.py

@@ -588,6 +588,30 @@ def name_version(self) -> SnapshotNameVersion:
         return SnapshotNameVersion(name=self.name, version=self.version)
 
 
+class MinimalSnapshot(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
+    without the overhead of parsing the full snapshot payload and fetching intervals.
+    """
+
+    name: str
+    version: str
+    dev_version_: t.Optional[str] = Field(alias="dev_version")
+    identifier: str
+    fingerprint: SnapshotFingerprint
+
+    @property
+    def snapshot_id(self) -> SnapshotId:
+        return SnapshotId(name=self.name, identifier=self.identifier)
+
+    @property
+    def name_version(self) -> SnapshotNameVersion:
+        return SnapshotNameVersion(name=self.name, version=self.version)
+
+    @property
+    def dev_version(self) -> str:
+        return self.dev_version_ or self.fingerprint.to_version()
+
+
 class Snapshot(PydanticModel, SnapshotInfoMixin):
     """A snapshot represents a node at a certain point in time.
 

sqlmesh/core/state_sync/base.py

@@ -23,6 +23,7 @@
     SnapshotTableCleanupTask,
     SnapshotTableInfo,
     SnapshotNameVersion,
+    MinimalSnapshot,
 )
 from sqlmesh.core.snapshot.definition import Interval, SnapshotIntervals
 from sqlmesh.utils import major_minor
@@ -98,21 +99,21 @@ def get_snapshots(
         """
 
     @abc.abstractmethod
-    def get_snapshot_ids_by_names(
+    def get_snapshots_by_names(
         self,
         snapshot_names: t.Iterable[str],
         current_ts: t.Optional[int] = None,
         exclude_expired: bool = True,
-    ) -> t.Set[SnapshotId]:
+    ) -> t.Set[MinimalSnapshot]:
         """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
+            snapshot_names: Iterable of snapshot names to fetch all snapshot 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
+            A set containing all the matched snapshot records. To fetch full snapshots, pass it into StateSync.get_snapshots()
         """
 
     @abc.abstractmethod

sqlmesh/core/state_sync/db/facade.py

@@ -29,6 +29,7 @@
 from sqlmesh.core.environment import Environment, EnvironmentStatements, EnvironmentSummary
 from sqlmesh.core.snapshot import (
     Snapshot,
+    MinimalSnapshot,
     SnapshotId,
     SnapshotIdLike,
     SnapshotInfoLike,
@@ -366,13 +367,13 @@ def get_snapshots(
         Snapshot.hydrate_with_intervals_by_version(snapshots.values(), intervals)
         return snapshots
 
-    def get_snapshot_ids_by_names(
+    def get_snapshots_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(
+    ) -> t.Set[MinimalSnapshot]:
+        return self.snapshot_state.get_snapshots_by_names(
             snapshot_names=snapshot_names, current_ts=current_ts, exclude_expired=exclude_expired
         )
 

sqlmesh/core/state_sync/db/snapshot.py

@@ -6,7 +6,6 @@
 from pathlib import Path
 from collections import defaultdict
 from sqlglot import exp
-from pydantic import Field
 
 from sqlmesh.core.engine_adapter import EngineAdapter
 from sqlmesh.core.state_sync.db.utils import (
@@ -27,12 +26,12 @@
     SnapshotNameVersion,
     SnapshotInfoLike,
     Snapshot,
+    MinimalSnapshot,
     SnapshotId,
     SnapshotFingerprint,
 )
 from sqlmesh.utils.migration import index_text_type, blob_text_type
 from sqlmesh.utils.date import now_timestamp, TimeLike, to_timestamp
-from sqlmesh.utils.pydantic import PydanticModel
 from sqlmesh.utils import unique
 
 if t.TYPE_CHECKING:
@@ -215,7 +214,7 @@ def _get_expired_snapshots(
             for snapshot in environment.snapshots
         }
 
-        def _is_snapshot_used(snapshot: SharedVersionSnapshot) -> bool:
+        def _is_snapshot_used(snapshot: MinimalSnapshot) -> bool:
             return (
                 snapshot.snapshot_id in promoted_snapshot_ids
                 or snapshot.snapshot_id not in expired_candidates
@@ -308,12 +307,12 @@ def get_snapshots(
         """
         return self._get_snapshots(snapshot_ids)
 
-    def get_snapshot_ids_by_names(
+    def get_snapshots_by_names(
         self,
         snapshot_names: t.Iterable[str],
         current_ts: t.Optional[int] = None,
         exclude_expired: bool = True,
-    ) -> t.Set[SnapshotId]:
+    ) -> t.Set[MinimalSnapshot]:
         """Return the snapshot id's for all versions of the specified snapshot names.
 
         Args:
@@ -334,14 +333,20 @@ def get_snapshot_ids_by_names(
             unexpired_expr = None
 
         return {
-            SnapshotId(name=name, identifier=identifier)
+            MinimalSnapshot(
+                name=name,
+                identifier=identifier,
+                version=version,
+                dev_version=dev_version,
+                fingerprint=SnapshotFingerprint.parse_raw(fingerprint),
+            )
             for where in snapshot_name_filter(
                 snapshot_names=snapshot_names,
                 batch_size=self.SNAPSHOT_BATCH_SIZE,
             )
-            for name, identifier in fetchall(
+            for name, identifier, version, dev_version, fingerprint in fetchall(
                 self.engine_adapter,
-                exp.select("name", "identifier")
+                exp.select("name", "identifier", "version", "dev_version", "fingerprint")
                 .from_(self.snapshots_table)
                 .where(where)
                 .and_(unexpired_expr),
@@ -631,7 +636,7 @@ def _get_snapshots_with_same_version(
         self,
         snapshots: t.Collection[SnapshotNameVersionLike],
         lock_for_update: bool = False,
-    ) -> t.List[SharedVersionSnapshot]:
+    ) -> t.List[MinimalSnapshot]:
         """Fetches all snapshots that share the same version as the snapshots.
 
         The output includes the snapshots with the specified identifiers.
@@ -668,7 +673,7 @@ def _get_snapshots_with_same_version(
             snapshot_rows.extend(fetchall(self.engine_adapter, query))
 
         return [
-            SharedVersionSnapshot(
+            MinimalSnapshot(
                 name=name,
                 identifier=identifier,
                 version=version,
@@ -751,23 +756,3 @@ def _auto_restatements_to_df(auto_restatements: t.Dict[SnapshotNameVersion, int]
             for name_version, ts in auto_restatements.items()
         ]
     )
-
-
-class SharedVersionSnapshot(PydanticModel):
-    """A stripped down version of a snapshot that is used for fetching snapshots that share the same version
-    with a significantly reduced parsing overhead.
-    """
-
-    name: str
-    version: str
-    dev_version_: t.Optional[str] = Field(alias="dev_version")
-    identifier: str
-    fingerprint: SnapshotFingerprint
-
-    @property
-    def snapshot_id(self) -> SnapshotId:
-        return SnapshotId(name=self.name, identifier=self.identifier)
-
-    @property
-    def dev_version(self) -> str:
-        return self.dev_version_ or self.fingerprint.to_version()

tests/core/state_sync/test_state_sync.py

@@ -3571,10 +3571,10 @@ def test_update_environment_statements(state_sync: EngineAdapterStateSync):
     ]
 
 
-def test_get_snapshot_ids_by_names(
+def test_get_snapshots_by_names(
     state_sync: EngineAdapterStateSync, make_snapshot: t.Callable[..., Snapshot]
 ):
-    assert state_sync.get_snapshot_ids_by_names(snapshot_names=[]) == set()
+    assert state_sync.get_snapshots_by_names(snapshot_names=[]) == set()
 
     snap_a_v1, snap_a_v2 = (
         make_snapshot(
@@ -3597,18 +3597,20 @@ def test_get_snapshot_ids_by_names(
 
     state_sync.push_snapshots([snap_a_v1, snap_a_v2, snap_b])
 
-    assert state_sync.get_snapshot_ids_by_names(snapshot_names=['"a"']) == {
+    assert {s.snapshot_id for s in state_sync.get_snapshots_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"']) == {
+    assert {
+        s.snapshot_id for s in state_sync.get_snapshots_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(
+def test_get_snapshots_by_names_include_expired(
     state_sync: EngineAdapterStateSync, make_snapshot: t.Callable[..., Snapshot]
 ):
     now_ts = now_timestamp()
@@ -3635,15 +3637,22 @@ def test_get_snapshot_ids_by_names_include_expired(
 
     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) == {
+    assert {
+        s.snapshot_id
+        for s in state_sync.get_snapshots_by_names(snapshot_names=['"a"'], current_ts=now_ts)
+    } == {normal_a.snapshot_id}
+    assert {
+        s.snapshot_id
+        for s in state_sync.get_snapshots_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}
+    assert {
+        s.snapshot_id
+        for s in state_sync.get_snapshots_by_names(
+            snapshot_names=['"a"'], current_ts=(now_ts - (10 * 1000))
+        )
+    } == {normal_a.snapshot_id, expired_a.snapshot_id}