fix: load materializations during cleanup

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:
@@ -1017,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
@@ -3182,6 +3188,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(
@@ -3373,9 +3381,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,

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,6 +88,121 @@ 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
+
+
+@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"
+    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)
+
+    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.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,
+    )
+
+    assert result.returncode == 0, result.stdout + result.stderr
+
+
 def test_janitor_cleanup_order(mocker: MockerFixture, tmp_path: Path):
     def setup_scenario():
         models_dir = tmp_path / "models"