chore: split up get snapshots operations (#2615)

Author: eakmanrq

sqlmesh/core/state_sync/engine_adapter.py

@@ -459,36 +459,24 @@ def _environments_query(
             return query.lock(copy=False)
         return query
 
-    def _get_snapshots(
+    def _get_snapshots_expressions(
         self,
         snapshot_ids: t.Optional[t.Iterable[SnapshotIdLike]] = None,
         lock_for_update: bool = False,
         hydrate_seeds: bool = False,
-        hydrate_intervals: bool = True,
-    ) -> t.Dict[SnapshotId, Snapshot]:
-        """Fetches specified snapshots or all snapshots.
-
-        Args:
-            snapshot_ids: The collection of snapshot like objects to fetch.
-            lock_for_update: Lock the snapshot rows for future update
-            hydrate_seeds: Whether to hydrate seed snapshots with the content.
-            hydrate_intervals: Whether to hydrate result snapshots with intervals.
-
-        Returns:
-            A dictionary of snapshot ids to snapshots for ones that could be found.
-        """
-        snapshots: t.Dict[SnapshotId, Snapshot] = {}
-        duplicates: t.Dict[SnapshotId, Snapshot] = {}
-        model_cache = ModelCache(self._context_path / c.CACHE)
-
+        batch_size: t.Optional[int] = None,
+    ) -> t.Iterator[exp.Expression]:
         for where in (
-            [None] if snapshot_ids is None else self._snapshot_id_filter(snapshot_ids, "snapshots")
+            [None]
+            if snapshot_ids is None
+            else self._snapshot_id_filter(snapshot_ids, alias="snapshots", batch_size=batch_size)
         ):
             query = (
                 exp.select(
                     "snapshots.snapshot",
                     "snapshots.name",
                     "snapshots.identifier",
+                    "snapshots.version",
                 )
                 .from_(exp.to_table(self.snapshots_table).as_("snapshots"))
                 .where(where)
@@ -504,17 +492,43 @@ def _get_snapshots(
                     ),
                     join_type="left",
                 )
-            elif lock_for_update:
+            else:
+                query = query.select(exp.Null().as_("content"))
+            if lock_for_update:
                 query = query.lock(copy=False)
+            yield query
 
-            for row in self._fetchall(query):
-                payload = json.loads(row[0])
+    def _get_snapshots(
+        self,
+        snapshot_ids: t.Optional[t.Iterable[SnapshotIdLike]] = None,
+        lock_for_update: bool = False,
+        hydrate_seeds: bool = False,
+        hydrate_intervals: bool = True,
+    ) -> t.Dict[SnapshotId, Snapshot]:
+        """Fetches specified snapshots or all snapshots.
+
+        Args:
+            snapshot_ids: The collection of snapshot like objects to fetch.
+            lock_for_update: Lock the snapshot rows for future update
+            hydrate_seeds: Whether to hydrate seed snapshots with the content.
+            hydrate_intervals: Whether to hydrate result snapshots with intervals.
 
-                def loader() -> Node:
-                    return parse_obj_as(Node, payload["node"])  # type: ignore
+        Returns:
+            A dictionary of snapshot ids to snapshots for ones that could be found.
+        """
+        snapshots: t.Dict[SnapshotId, Snapshot] = {}
+        duplicates: t.Dict[SnapshotId, Snapshot] = {}
+        model_cache = ModelCache(self._context_path / c.CACHE)
 
-                payload["node"] = model_cache.get_or_load(f"{row[1]}_{row[2]}", loader=loader)  # type: ignore
-                snapshot = Snapshot(**payload)
+        for query in self._get_snapshots_expressions(snapshot_ids, lock_for_update, hydrate_seeds):
+            for serialized_snapshot, name, identifier, _, seed_content in self._fetchall(query):
+                snapshot = parse_snapshot(
+                    model_cache,
+                    serialized_snapshot=serialized_snapshot,
+                    name=name,
+                    identifier=identifier,
+                    seed_content=seed_content,
+                )
                 snapshot_id = snapshot.snapshot_id
                 if snapshot_id in snapshots:
                     other = duplicates.get(snapshot_id, snapshots[snapshot_id])
@@ -525,9 +539,6 @@ def loader() -> Node:
                 else:
                     snapshots[snapshot_id] = snapshot
 
-                if hydrate_seeds and isinstance(snapshot.node, SeedModel) and row[-1]:
-                    snapshot.node = t.cast(SeedModel, snapshot.node).to_hydrated(row[-1])
-
         if snapshots and hydrate_intervals:
             _, intervals = self._get_snapshot_intervals(snapshots.values())
             Snapshot.hydrate_with_intervals_by_version(snapshots.values(), intervals)
@@ -1452,6 +1463,27 @@ def _snapshot_to_json(snapshot: Snapshot) -> str:
     return snapshot.json(exclude={"intervals", "dev_intervals"})
 
 
+def parse_snapshot(
+    model_cache: ModelCache,
+    serialized_snapshot: str,
+    name: str,
+    identifier: str,
+    seed_content: t.Optional[str],
+) -> Snapshot:
+    payload = json.loads(serialized_snapshot)
+
+    def loader() -> Node:
+        return parse_obj_as(Node, payload["node"])  # type: ignore
+
+    payload["node"] = model_cache.get_or_load(f"{name}_{identifier}", loader=loader)  # type: ignore
+    snapshot = Snapshot(**payload)
+
+    if seed_content and isinstance(snapshot.node, SeedModel):
+        snapshot.node = snapshot.node.to_hydrated(seed_content)
+
+    return snapshot
+
+
 class LazilyParsedSnapshots:
     def __init__(self, raw_snapshots: t.Dict[SnapshotId, t.Dict[str, t.Any]]):
         self._raw_snapshots = raw_snapshots

tests/core/test_state_sync.py

@@ -2193,13 +2193,17 @@ def test_snapshot_batching(state_sync, mocker, make_snapshot):
                 ).json(),
                 "a",
                 "1",
+                "1",
+                None,
             ],
             [
                 make_snapshot(
                     SqlModel(name="a", query=parse_one("select 2")),
                 ).json(),
                 "a",
                 "2",
+                "2",
+                None,
             ],
         ],
         [
@@ -2209,6 +2213,8 @@ def test_snapshot_batching(state_sync, mocker, make_snapshot):
                 ).json(),
                 "a",
                 "3",
+                "3",
+                None,
             ],
         ],
     ]