From ae7439be53a0d025d2d3b16f0b5c68cca3074893 Mon Sep 17 00:00:00 2001
From: lafirm <136463254+lafirm@users.noreply.github.com>
Date: Mon, 26 Oct 2026 10:28:29 +0530
Subject: [PATCH 1/6] fix: load materializations during cleanup

Signed-off-by: lafirm <136463254+lafirm@users.noreply.github.com>
---
 sqlmesh/core/context.py                     |  2 ++
 tests/core/integration/test_aux_commands.py | 23 +++++++++++++++++++++
 2 files changed, 25 insertions(+)

diff --git a/sqlmesh/core/context.py b/sqlmesh/core/context.py
index 39590a0fa3..fb341c4fa6 100644
--- a/sqlmesh/core/context.py
+++ b/sqlmesh/core/context.py
@@ -3182,6 +3182,8 @@ def _run_janitor(
         current_ts = now_timestamp()
         failures: t.List[str] = []
 
+        self._load_materializations()
+
         # Clean up expired environments by removing their views and schemas
         failures.extend(
             self._cleanup_environments(
diff --git a/tests/core/integration/test_aux_commands.py b/tests/core/integration/test_aux_commands.py
index 7de585576d..c91fa7a30c 100644
--- a/tests/core/integration/test_aux_commands.py
+++ b/tests/core/integration/test_aux_commands.py
@@ -86,6 +86,29 @@ def test_table_name(init_and_plan_context: t.Callable):
     )
 
 
+def test_janitor_loads_materialization(mocker: MockerFixture, tmp_path: Path):
+    models_dir = tmp_path / "models"
+    models_dir.mkdir(exist_ok=True)
+
+    test_model = models_dir / "model.sql"
+
+    with open(test_model, "w") as f:
+        f.write(
+            "MODEL(name test.model, kind CUSTOM (materialization 'custom_full')); SELECT 1 AS col"
+        )
+
+    config = Config(model_defaults=ModelDefaultsConfig(dialect="duckdb"))
+    ctx = Context(paths=[tmp_path], config=config)
+    ctx.plan("dev", no_prompts=True, auto_apply=True)
+
+    load_materializations_spy = mocker.spy(Context, "_load_materializations")
+
+    ctx.invalidate_environment("dev")
+    ctx.run_janitor(ignore_ttl=True)
+
+    load_materializations_spy.assert_called_once()
+
+
 def test_janitor_cleanup_order(mocker: MockerFixture, tmp_path: Path):
     def setup_scenario():
         models_dir = tmp_path / "models"

From a6ca44cf7eec50ff416836e7da7aa4a84371650d Mon Sep 17 00:00:00 2001
From: lafirm <136463254+lafirm@users.noreply.github.com>
Date: Wed, 4 Nov 2026 11:29:50 +0530
Subject: [PATCH 2/6] fix: add self._materializations_loaded flag in ctxt

Signed-off-by: lafirm <136463254+lafirm@users.noreply.github.com>
---
 sqlmesh/core/context.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/sqlmesh/core/context.py b/sqlmesh/core/context.py
index fb341c4fa6..dbe5cd82e5 100644
--- a/sqlmesh/core/context.py
+++ b/sqlmesh/core/context.py
@@ -420,6 +420,7 @@ def __init__(
         self._engine_adapter: t.Optional[EngineAdapter] = None
         self._linters: t.Dict[str, Linter] = {}
         self._loaded: bool = False
+        self._materializations_loaded: bool = False
         self._load_state: bool = load_state
         self._uncached_model_names: t.Set[str] = set()
         self._selector_cls = selector or NativeSelector
@@ -807,6 +808,7 @@ def load(
         )
 
         self._loaded = True
+        self._materializations_loaded = True
         return self
 
     def _update_model_schemas_and_validate(self, model_fqns: t.Optional[t.Set[str]] = None) -> None:
@@ -3375,9 +3377,10 @@ def _register_notification_targets(self) -> None:
         )
 
     def _load_materializations(self) -> None:
-        if not self._loaded:
+        if not self._materializations_loaded:
             for loader in self._loaders:
                 loader.load_materializations()
+            self._materializations_loaded = True
 
     def _select_models_for_run(
         self,

From 846271e470e3574c0be0a7c889a938347c73a2d7 Mon Sep 17 00:00:00 2001
From: lafirm <136463254+lafirm@users.noreply.github.com>
Date: Wed, 4 Nov 2026 12:22:08 +0530
Subject: [PATCH 3/6] fix: modify janitor test to implement a new sub-porcess

Signed-off-by: lafirm <136463254+lafirm@users.noreply.github.com>
---
 tests/core/integration/test_aux_commands.py | 86 +++++++++++++++++----
 1 file changed, 69 insertions(+), 17 deletions(-)

diff --git a/tests/core/integration/test_aux_commands.py b/tests/core/integration/test_aux_commands.py
index c91fa7a30c..0c2f2a952a 100644
--- a/tests/core/integration/test_aux_commands.py
+++ b/tests/core/integration/test_aux_commands.py
@@ -1,6 +1,8 @@
 from __future__ import annotations
 
 import typing as t
+import subprocess
+import sys
 from unittest.mock import patch
 import pytest
 from pathlib import Path
@@ -86,27 +88,77 @@ def test_table_name(init_and_plan_context: t.Callable):
     )
 
 
-def test_janitor_loads_materialization(mocker: MockerFixture, tmp_path: Path):
-    models_dir = tmp_path / "models"
-    models_dir.mkdir(exist_ok=True)
-
-    test_model = models_dir / "model.sql"
-
-    with open(test_model, "w") as f:
-        f.write(
-            "MODEL(name test.model, kind CUSTOM (materialization 'custom_full')); SELECT 1 AS col"
-        )
-
-    config = Config(model_defaults=ModelDefaultsConfig(dialect="duckdb"))
-    ctx = Context(paths=[tmp_path], config=config)
+def test_janitor_loads_materialization(tmp_path: Path):
+    config_file_path = Path("config.yaml")
+    database = tmp_path / "janitor_test.duckdb"
+    materialization_name = "dummy_materialization"
+    materialization_file_path = Path(f"materializations/{materialization_name}.py")
+    schema_name = "test"
+    model_name = "dummy_model"
+    model_file_path = Path(f"models/{model_name}.sql")
+    config_yaml = f"""\
+gateways:
+    local:
+        connection:
+            type: duckdb
+            database: {database}
+default_gateway: local
+model_defaults:
+    dialect: duckdb
+"""
+    materialization_code = f"""from sqlmesh.core.snapshot.evaluator import CustomMaterialization
+class DummyMaterialization(CustomMaterialization):
+    NAME = "{materialization_name}"
+    def insert(self, table_name, query_or_df, model, is_first_insert, render_kwargs, **kwargs):
+        self._replace_query_for_model(model, table_name, query_or_df, render_kwargs)
+"""
+    model_code = f"""MODEL(name {schema_name}.{model_name}, kind CUSTOM (materialization '{materialization_name}'));
+SELECT 1 AS col"""
+
+    create_temp_file(tmp_path, config_file_path, config_yaml)
+    create_temp_file(tmp_path, materialization_file_path, materialization_code)
+    create_temp_file(tmp_path, model_file_path, model_code)
+
+    ctx = Context(paths=[tmp_path])
     ctx.plan("dev", no_prompts=True, auto_apply=True)
 
-    load_materializations_spy = mocker.spy(Context, "_load_materializations")
-
+    snapshot = ctx.get_snapshot(f"{schema_name}.{model_name}")
+    assert snapshot
+    table_name = snapshot.table_name()
+    assert ctx.engine_adapter.table_exists(table_name)
+    assert ctx.state_sync.get_snapshots([snapshot.snapshot_id])
+    assert ctx._loaded
+    assert ctx._materializations_loaded
     ctx.invalidate_environment("dev")
-    ctx.run_janitor(ignore_ttl=True)
+    ctx.close()
+
+    # A fresh interpreter cannot reuse materialization classes loaded during setup.
+    program = """\
+import sys
+from sqlmesh.core.context import Context
+
+ctx = Context(paths=[sys.argv[1]], load=False)
+table_name = sys.argv[2]
+assert not ctx._loaded
+assert not ctx._materializations_loaded
+assert ctx.engine_adapter.table_exists(table_name)
+
+assert ctx.run_janitor(ignore_ttl=True)
+
+assert not ctx.engine_adapter.table_exists(table_name)
+assert not ctx._loaded
+assert ctx._materializations_loaded
+ctx.close()
+"""
+    result = subprocess.run(
+        [sys.executable, "-c", program, str(tmp_path), table_name],
+        cwd=tmp_path,
+        capture_output=True,
+        text=True,
+        timeout=120,
+    )
 
-    load_materializations_spy.assert_called_once()
+    assert result.returncode == 0, result.stdout + result.stderr
 
 
 def test_janitor_cleanup_order(mocker: MockerFixture, tmp_path: Path):

From c5ad0fbd41acc452c522d1157c03975a3e2ed542 Mon Sep 17 00:00:00 2001
From: lafirm <136463254+lafirm@users.noreply.github.com>
Date: Wed, 4 Nov 2026 15:42:26 +0530
Subject: [PATCH 4/6] fix: add test for run loads materializations only once

Signed-off-by: lafirm <136463254+lafirm@users.noreply.github.com>
---
 tests/core/integration/test_aux_commands.py | 22 +++++++++++++++++++++
 1 file changed, 22 insertions(+)

diff --git a/tests/core/integration/test_aux_commands.py b/tests/core/integration/test_aux_commands.py
index 0c2f2a952a..5026c0fca6 100644
--- a/tests/core/integration/test_aux_commands.py
+++ b/tests/core/integration/test_aux_commands.py
@@ -88,6 +88,28 @@ def test_table_name(init_and_plan_context: t.Callable):
     )
 
 
+@time_machine.travel("2023-01-08 15:00:00 UTC")
+def test_run_loads_materializations_once(init_and_plan_context: t.Callable, mocker: MockerFixture):
+    """Verify janitor and run dispatch loading while loaders execute only once."""
+    context, plan = init_and_plan_context("examples/sushi")
+    context.apply(plan)
+    load_dispatch_spy = mocker.spy(context, "_load_materializations")
+    loader_spies = [mocker.spy(loader, "load_materializations") for loader in context._loaders]
+    assert loader_spies
+    assert context._materializations_loaded is True
+
+    # Reset the guard to exercise lazy loading through run's janitor and scheduler paths.
+    context._materializations_loaded = False
+    context.run(skip_janitor=False)
+
+    # The janitor and run each call the context wrapper; the second call is guarded.
+    assert load_dispatch_spy.call_count == 2
+
+    # Only the janitor's first dispatch reaches the loaders.
+    for loader_spy in loader_spies:
+        assert loader_spy.call_count == 1
+
+
 def test_janitor_loads_materialization(tmp_path: Path):
     config_file_path = Path("config.yaml")
     database = tmp_path / "janitor_test.duckdb"

From 8a7edf808eed35ffc5ef35d3bdedd6a70a5a285a Mon Sep 17 00:00:00 2001
From: lafirm <136463254+lafirm@users.noreply.github.com>
Date: Thu, 5 Nov 2026 11:04:56 +0530
Subject: [PATCH 5/6] fix: load materializations before destroy invalidates
 environments

Signed-off-by: lafirm <136463254+lafirm@users.noreply.github.com>
---
 sqlmesh/core/context.py                     |  3 +++
 tests/core/integration/test_aux_commands.py | 20 ++++++++++++++++++++
 2 files changed, 23 insertions(+)

diff --git a/sqlmesh/core/context.py b/sqlmesh/core/context.py
index dbe5cd82e5..d7adcaa7f5 100644
--- a/sqlmesh/core/context.py
+++ b/sqlmesh/core/context.py
@@ -3158,6 +3158,9 @@ def _context_diff(
         )
 
     def _destroy(self) -> bool:
+        # Load materializations before any state writes so a failure aborts destroy without invalidating environments
+        self._load_materializations()
+
         # Invalidate all environments, including prod
         for environment in self.state_reader.get_environments():
             self.state_sync.invalidate_environment(name=environment.name, protect_prod=False)
diff --git a/tests/core/integration/test_aux_commands.py b/tests/core/integration/test_aux_commands.py
index 5026c0fca6..16f5fcaad2 100644
--- a/tests/core/integration/test_aux_commands.py
+++ b/tests/core/integration/test_aux_commands.py
@@ -110,6 +110,26 @@ def test_run_loads_materializations_once(init_and_plan_context: t.Callable, mock
         assert loader_spy.call_count == 1
 
 
+@time_machine.travel("2023-01-08 15:00:00 UTC")
+def test_destroy_fails_before_invalidating_environments_on_materialization_error(
+    init_and_plan_context: t.Callable, mocker: MockerFixture
+):
+    context, plan = init_and_plan_context("examples/sushi")
+    context.apply(plan)
+    assert context.state_reader.get_environment("prod").expiration_ts is None
+
+    for loader in context._loaders:
+        mocker.patch.object(
+            loader, "load_materializations", side_effect=SQLMeshError("broken materialization")
+        )
+    context._materializations_loaded = False
+
+    with pytest.raises(SQLMeshError, match="broken materialization"):
+        context._destroy()
+
+    assert context.state_reader.get_environment("prod").expiration_ts is None
+
+
 def test_janitor_loads_materialization(tmp_path: Path):
     config_file_path = Path("config.yaml")
     database = tmp_path / "janitor_test.duckdb"

From 241774c14f355bab3d438f3b0b1f8be88b07aabc Mon Sep 17 00:00:00 2001
From: lafirm <136463254+lafirm@users.noreply.github.com>
Date: Thu, 5 Nov 2026 13:16:20 +0530
Subject: [PATCH 6/6] fix: load materializations before destroy snapshot reads

Signed-off-by: lafirm <136463254+lafirm@users.noreply.github.com>
---
 sqlmesh/core/context.py                     | 7 ++++---
 tests/core/integration/test_aux_commands.py | 2 +-
 2 files changed, 5 insertions(+), 4 deletions(-)

diff --git a/sqlmesh/core/context.py b/sqlmesh/core/context.py
index d7adcaa7f5..fd7f8ad071 100644
--- a/sqlmesh/core/context.py
+++ b/sqlmesh/core/context.py
@@ -1019,6 +1019,10 @@ def run_janitor(
 
     @python_api_analytics
     def destroy(self) -> bool:
+        # Load first so snapshot unpickling cannot import a custom-kind module that this load then reimports.
+        # If loading fails, abort before invalidating any environments.
+        self._load_materializations()
+
         success = False
 
         # Collect resources to be deleted
@@ -3158,9 +3162,6 @@ def _context_diff(
         )
 
     def _destroy(self) -> bool:
-        # Load materializations before any state writes so a failure aborts destroy without invalidating environments
-        self._load_materializations()
-
         # Invalidate all environments, including prod
         for environment in self.state_reader.get_environments():
             self.state_sync.invalidate_environment(name=environment.name, protect_prod=False)
diff --git a/tests/core/integration/test_aux_commands.py b/tests/core/integration/test_aux_commands.py
index 16f5fcaad2..070f38d780 100644
--- a/tests/core/integration/test_aux_commands.py
+++ b/tests/core/integration/test_aux_commands.py
@@ -125,7 +125,7 @@ def test_destroy_fails_before_invalidating_environments_on_materialization_error
     context._materializations_loaded = False
 
     with pytest.raises(SQLMeshError, match="broken materialization"):
-        context._destroy()
+        context.destroy()
 
     assert context.state_reader.get_environment("prod").expiration_ts is None