Skip to content

Commit 7975a8d

Browse files
lafirmmday-io
andauthored
fix: load materializations during cleanup (#6162)
Signed-off-by: lafirm <136463254+lafirm@users.noreply.github.com> Co-authored-by: Michael Day <mdaytn@gmail.com>
1 parent f0ed98e commit 7975a8d

2 files changed

Lines changed: 127 additions & 1 deletion

File tree

‎sqlmesh/core/context.py‎

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -420,6 +420,7 @@ def __init__(
420420
self._engine_adapter: t.Optional[EngineAdapter] = None
421421
self._linters: t.Dict[str, Linter] = {}
422422
self._loaded: bool = False
423+
self._materializations_loaded: bool = False
423424
self._load_state: bool = load_state
424425
self._uncached_model_names: t.Set[str] = set()
425426
self._selector_cls = selector or NativeSelector
@@ -807,6 +808,7 @@ def load(
807808
)
808809

809810
self._loaded = True
811+
self._materializations_loaded = True
810812
return self
811813

812814
def _update_model_schemas_and_validate(self, model_fqns: t.Optional[t.Set[str]] = None) -> None:
@@ -1017,6 +1019,10 @@ def run_janitor(
10171019

10181020
@python_api_analytics
10191021
def destroy(self) -> bool:
1022+
# Load first so snapshot unpickling cannot import a custom-kind module that this load then reimports.
1023+
# If loading fails, abort before invalidating any environments.
1024+
self._load_materializations()
1025+
10201026
success = False
10211027

10221028
# Collect resources to be deleted
@@ -3182,6 +3188,8 @@ def _run_janitor(
31823188
current_ts = now_timestamp()
31833189
failures: t.List[str] = []
31843190

3191+
self._load_materializations()
3192+
31853193
# Clean up expired environments by removing their views and schemas
31863194
failures.extend(
31873195
self._cleanup_environments(
@@ -3373,9 +3381,10 @@ def _register_notification_targets(self) -> None:
33733381
)
33743382

33753383
def _load_materializations(self) -> None:
3376-
if not self._loaded:
3384+
if not self._materializations_loaded:
33773385
for loader in self._loaders:
33783386
loader.load_materializations()
3387+
self._materializations_loaded = True
33793388

33803389
def _select_models_for_run(
33813390
self,

‎tests/core/integration/test_aux_commands.py‎

Lines changed: 117 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,8 @@
11
from __future__ import annotations
22

33
import typing as t
4+
import subprocess
5+
import sys
46
from unittest.mock import patch
57
import pytest
68
from pathlib import Path
@@ -86,6 +88,121 @@ def test_table_name(init_and_plan_context: t.Callable):
8688
)
8789

8890

91+
@time_machine.travel("2023-01-08 15:00:00 UTC")
92+
def test_run_loads_materializations_once(init_and_plan_context: t.Callable, mocker: MockerFixture):
93+
"""Verify janitor and run dispatch loading while loaders execute only once."""
94+
context, plan = init_and_plan_context("examples/sushi")
95+
context.apply(plan)
96+
load_dispatch_spy = mocker.spy(context, "_load_materializations")
97+
loader_spies = [mocker.spy(loader, "load_materializations") for loader in context._loaders]
98+
assert loader_spies
99+
assert context._materializations_loaded is True
100+
101+
# Reset the guard to exercise lazy loading through run's janitor and scheduler paths.
102+
context._materializations_loaded = False
103+
context.run(skip_janitor=False)
104+
105+
# The janitor and run each call the context wrapper; the second call is guarded.
106+
assert load_dispatch_spy.call_count == 2
107+
108+
# Only the janitor's first dispatch reaches the loaders.
109+
for loader_spy in loader_spies:
110+
assert loader_spy.call_count == 1
111+
112+
113+
@time_machine.travel("2023-01-08 15:00:00 UTC")
114+
def test_destroy_fails_before_invalidating_environments_on_materialization_error(
115+
init_and_plan_context: t.Callable, mocker: MockerFixture
116+
):
117+
context, plan = init_and_plan_context("examples/sushi")
118+
context.apply(plan)
119+
assert context.state_reader.get_environment("prod").expiration_ts is None
120+
121+
for loader in context._loaders:
122+
mocker.patch.object(
123+
loader, "load_materializations", side_effect=SQLMeshError("broken materialization")
124+
)
125+
context._materializations_loaded = False
126+
127+
with pytest.raises(SQLMeshError, match="broken materialization"):
128+
context.destroy()
129+
130+
assert context.state_reader.get_environment("prod").expiration_ts is None
131+
132+
133+
def test_janitor_loads_materialization(tmp_path: Path):
134+
config_file_path = Path("config.yaml")
135+
database = tmp_path / "janitor_test.duckdb"
136+
materialization_name = "dummy_materialization"
137+
materialization_file_path = Path(f"materializations/{materialization_name}.py")
138+
schema_name = "test"
139+
model_name = "dummy_model"
140+
model_file_path = Path(f"models/{model_name}.sql")
141+
config_yaml = f"""\
142+
gateways:
143+
local:
144+
connection:
145+
type: duckdb
146+
database: {database}
147+
default_gateway: local
148+
model_defaults:
149+
dialect: duckdb
150+
"""
151+
materialization_code = f"""from sqlmesh.core.snapshot.evaluator import CustomMaterialization
152+
class DummyMaterialization(CustomMaterialization):
153+
NAME = "{materialization_name}"
154+
def insert(self, table_name, query_or_df, model, is_first_insert, render_kwargs, **kwargs):
155+
self._replace_query_for_model(model, table_name, query_or_df, render_kwargs)
156+
"""
157+
model_code = f"""MODEL(name {schema_name}.{model_name}, kind CUSTOM (materialization '{materialization_name}'));
158+
SELECT 1 AS col"""
159+
160+
create_temp_file(tmp_path, config_file_path, config_yaml)
161+
create_temp_file(tmp_path, materialization_file_path, materialization_code)
162+
create_temp_file(tmp_path, model_file_path, model_code)
163+
164+
ctx = Context(paths=[tmp_path])
165+
ctx.plan("dev", no_prompts=True, auto_apply=True)
166+
167+
snapshot = ctx.get_snapshot(f"{schema_name}.{model_name}")
168+
assert snapshot
169+
table_name = snapshot.table_name()
170+
assert ctx.engine_adapter.table_exists(table_name)
171+
assert ctx.state_sync.get_snapshots([snapshot.snapshot_id])
172+
assert ctx._loaded
173+
assert ctx._materializations_loaded
174+
ctx.invalidate_environment("dev")
175+
ctx.close()
176+
177+
# A fresh interpreter cannot reuse materialization classes loaded during setup.
178+
program = """\
179+
import sys
180+
from sqlmesh.core.context import Context
181+
182+
ctx = Context(paths=[sys.argv[1]], load=False)
183+
table_name = sys.argv[2]
184+
assert not ctx._loaded
185+
assert not ctx._materializations_loaded
186+
assert ctx.engine_adapter.table_exists(table_name)
187+
188+
assert ctx.run_janitor(ignore_ttl=True)
189+
190+
assert not ctx.engine_adapter.table_exists(table_name)
191+
assert not ctx._loaded
192+
assert ctx._materializations_loaded
193+
ctx.close()
194+
"""
195+
result = subprocess.run(
196+
[sys.executable, "-c", program, str(tmp_path), table_name],
197+
cwd=tmp_path,
198+
capture_output=True,
199+
text=True,
200+
timeout=120,
201+
)
202+
203+
assert result.returncode == 0, result.stdout + result.stderr
204+
205+
89206
def test_janitor_cleanup_order(mocker: MockerFixture, tmp_path: Path):
90207
def setup_scenario():
91208
models_dir = tmp_path / "models"

0 commit comments

Comments
 (0)