diff --git a/sqlmesh/core/context.py b/sqlmesh/core/context.py index eca42e1295..86c3cd9043 100644 --- a/sqlmesh/core/context.py +++ b/sqlmesh/core/context.py @@ -745,7 +745,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: diff --git a/tests/core/test_context.py b/tests/core/test_context.py index 75737f1edb..658d3659b6 100644 --- a/tests/core/test_context.py +++ b/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()