perf: selective snapshot hydration (#6016)

Signed-off-by: Andreas Fredhøi <andreas.fredhoi@fresio.no>
Co-authored-by: Cortland Goffena <30168413+cmgoffena13@users.noreply.github.com>

Author: fresioAS

sqlmesh/core/context.py

@@ -748,7 +748,25 @@ def load(
             prod = self.state_reader.get_environment(c.PROD)
 
             if prod:
-                for snapshot in self.state_reader.get_snapshots(prod.snapshots).values():
+                # Environment snapshot infos already contain the names we need to distinguish
+                # local nodes from nodes owned by another project. Hydrating local snapshots
+                # here is wasteful: their payloads are not used, and a large remote state can
+                # spend most of Context.load() decoding those model object graphs.
+                #
+                # Only hydrate names which are absent locally. These can be deleted nodes from
+                # this project or nodes from another project which must be merged into the
+                # context. The project field on the hydrated node disambiguates the two cases.
+                remote_snapshot_infos = []
+                for snapshot_info in prod.snapshots:
+                    local_store = (
+                        self._standalone_audits if snapshot_info.is_audit else self._models
+                    )
+                    if snapshot_info.name in local_store:
+                        uncached.add(snapshot_info.name)
+                    else:
+                        remote_snapshot_infos.append(snapshot_info)
+
+                for snapshot in self.state_reader.get_snapshots(remote_snapshot_infos).values():
                     if snapshot.node.project in self._projects:
                         uncached.add(snapshot.name)
                     else:

tests/core/test_context.py

@@ -17,6 +17,7 @@
 
 import sqlmesh.core.constants
 from sqlmesh.cli.project_init import init_example_project
+from sqlmesh.core.audit import StandaloneAudit
 from sqlmesh.core.console import TerminalConsole
 from sqlmesh.core import dialect as d, constants as c
 from sqlmesh.core.config import (
@@ -276,6 +277,78 @@ def test_render_seed_model(sushi_context, assert_exp_eq):
     )
 
 
+@pytest.mark.slow
+def test_load_only_hydrates_remote_snapshots_missing_locally(sushi_context: Context) -> None:
+    prod = sushi_context.state_reader.get_environment("prod")
+    assert prod is not None
+    sushi_context._projects = {"local_project"}
+
+    local_node_names = {*sushi_context.models, *sushi_context.standalone_audits}
+    expected_remote_names = {
+        snapshot_info.name
+        for snapshot_info in prod.snapshots
+        if snapshot_info.name not in local_node_names
+    }
+    assert len(expected_remote_names) < len(prod.snapshots)
+
+    with patch.object(
+        sushi_context.state_reader,
+        "get_snapshots",
+        wraps=sushi_context.state_reader.get_snapshots,
+    ) as get_snapshots_mock:
+        sushi_context.load(update_schemas=False)
+
+    load_snapshot_names = {
+        snapshot_info.name
+        for call_args in get_snapshots_mock.call_args_list
+        for snapshot_info in call_args.args[0]
+    }
+    assert load_snapshot_names == expected_remote_names
+
+
+@pytest.mark.slow
+def test_load_hydrates_opposite_type_snapshot_name_collision(
+    sushi_context: Context, make_snapshot: t.Callable
+) -> None:
+    prod = sushi_context.state_reader.get_environment("prod")
+    assert prod is not None
+    sushi_context._projects = {"local_project"}
+
+    local_model = next(iter(sushi_context.models.values()))
+    remote_audit_snapshot = make_snapshot(
+        StandaloneAudit(
+            name=local_model.fqn,
+            query=parse_one("SELECT NULL LIMIT 0"),
+            project="remote_project",
+        )
+    )
+    remote_audit_snapshot.categorize_as(SnapshotChangeCategory.BREAKING)
+    prod_with_collision = prod.copy(
+        update={"snapshots_": [*prod.snapshots, remote_audit_snapshot.table_info]}
+    )
+
+    with (
+        patch.object(
+            sushi_context.state_reader,
+            "get_environment",
+            return_value=prod_with_collision,
+        ),
+        patch.object(
+            sushi_context.state_reader,
+            "get_snapshots",
+            wraps=sushi_context.state_reader.get_snapshots,
+        ) as get_snapshots_mock,
+    ):
+        sushi_context.load(update_schemas=False)
+
+    requested_snapshot_names = {
+        snapshot_info.name
+        for call_args in get_snapshots_mock.call_args_list
+        for snapshot_info in call_args.args[0]
+    }
+    assert remote_audit_snapshot.name in requested_snapshot_names
+
+
 @pytest.mark.slow
 def test_diff(sushi_context: Context, mocker: MockerFixture):
     mock_console = mocker.Mock()