diff --git a/sqlmesh/core/context.py b/sqlmesh/core/context.py index 39590a0fa3..fd7f8ad071 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: @@ -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, diff --git a/tests/core/integration/test_aux_commands.py b/tests/core/integration/test_aux_commands.py index 7de585576d..070f38d780 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,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"