From 6f0d3639f536aa15c9cea62b30d5adf624517945 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Andreas=20Fredh=C3=B8i?= Date: Wed, 2 Sep 2026 01:24:22 +0200 Subject: [PATCH 1/2] perf: avoid hydrating local snapshots during context load MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Andreas Fredhøi --- sqlmesh/core/context.py | 18 +++++++++++++++++- tests/core/test_context.py | 28 ++++++++++++++++++++++++++++ 2 files changed, 45 insertions(+), 1 deletion(-) diff --git a/sqlmesh/core/context.py b/sqlmesh/core/context.py index eca42e1295..2a304c88be 100644 --- a/sqlmesh/core/context.py +++ b/sqlmesh/core/context.py @@ -745,7 +745,23 @@ 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. + local_node_names = {*self._models, *self._standalone_audits} + remote_snapshot_infos = [] + for snapshot_info in prod.snapshots: + if snapshot_info.name in local_node_names: + 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..4cd545c647 100644 --- a/tests/core/test_context.py +++ b/tests/core/test_context.py @@ -276,6 +276,34 @@ 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 + + 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_diff(sushi_context: Context, mocker: MockerFixture): mock_console = mocker.Mock() From 2167c12b4dd5b82488c0b440d410141a3704947b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Andreas=20Fredh=C3=B8i?= Date: Wed, 2 Sep 2026 01:38:44 +0200 Subject: [PATCH 2/2] fix: preserve snapshot node type collision handling MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Andreas Fredhøi --- sqlmesh/core/context.py | 6 +++-- tests/core/test_context.py | 45 ++++++++++++++++++++++++++++++++++++++ 2 files changed, 49 insertions(+), 2 deletions(-) diff --git a/sqlmesh/core/context.py b/sqlmesh/core/context.py index 2a304c88be..86c3cd9043 100644 --- a/sqlmesh/core/context.py +++ b/sqlmesh/core/context.py @@ -753,10 +753,12 @@ def load( # 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. - local_node_names = {*self._models, *self._standalone_audits} remote_snapshot_infos = [] for snapshot_info in prod.snapshots: - if snapshot_info.name in local_node_names: + 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) diff --git a/tests/core/test_context.py b/tests/core/test_context.py index 4cd545c647..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 ( @@ -280,6 +281,7 @@ def test_render_seed_model(sushi_context, assert_exp_eq): 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 = { @@ -304,6 +306,49 @@ def test_load_only_hydrates_remote_snapshots_missing_locally(sushi_context: Cont 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()