Skip to content

Commit b16ceb6

Browse files
Chore: Add integration test to verify before all execution order
1 parent a7feeb7 commit b16ceb6

File tree

1 file changed

+70
-1
lines changed

1 file changed

+70
-1
lines changed

tests/core/test_integration.py

Lines changed: 70 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -42,7 +42,7 @@
4242
from sqlmesh.core.context import Context
4343
from sqlmesh.core.config.categorizer import CategorizerConfig
4444
from sqlmesh.core.config.plan import PlanConfig
45-
from sqlmesh.core.engine_adapter import EngineAdapter
45+
from sqlmesh.core.engine_adapter import EngineAdapter, DuckDBEngineAdapter
4646
from sqlmesh.core.environment import EnvironmentNamingInfo
4747
from sqlmesh.core.macros import macro
4848
from sqlmesh.core.model import (
@@ -6427,6 +6427,75 @@ def test_environment_statements_error_handling(tmp_path: Path):
64276427
ctx.plan(auto_apply=True, no_prompts=True)
64286428

64296429

6430+
def test_before_all_after_all_execution_order(tmp_path: Path, mocker: MockerFixture):
6431+
model = """
6432+
MODEL (
6433+
name test_schema.model_that_depends_on_before_all,
6434+
kind FULL,
6435+
);
6436+
6437+
SELECT id, value FROM before_all_created_table
6438+
"""
6439+
6440+
models_dir = tmp_path / "models"
6441+
models_dir.mkdir()
6442+
6443+
with open(models_dir / "model.sql", "w") as f:
6444+
f.write(model)
6445+
6446+
# before_all statement that creates a table that the above model depends on
6447+
before_all_statement = (
6448+
"CREATE TABLE IF NOT EXISTS before_all_created_table AS SELECT 1 AS id, 'test' AS value"
6449+
)
6450+
6451+
# after_all that depends on the model
6452+
after_all_statement = "CREATE TABLE IF NOT EXISTS after_all_created_table AS SELECT id, value FROM test_schema.model_that_depends_on_before_all"
6453+
6454+
config = Config(
6455+
model_defaults=ModelDefaultsConfig(dialect="duckdb"),
6456+
before_all=[before_all_statement],
6457+
after_all=[after_all_statement],
6458+
)
6459+
6460+
execute_calls: t.List[str] = []
6461+
6462+
original_duckdb_execute = DuckDBEngineAdapter.execute
6463+
6464+
def track_duckdb_execute(self, expression, **kwargs):
6465+
sql = expression if isinstance(expression, str) else expression.sql(dialect="duckdb")
6466+
state_tables = [
6467+
"_snapshots",
6468+
"_environments",
6469+
"_versions",
6470+
"_intervals",
6471+
"_auto_restatements",
6472+
"_environment_statements",
6473+
"_plan_dags",
6474+
]
6475+
6476+
# to ignore the state queries
6477+
if not any(table in sql.lower() for table in state_tables):
6478+
execute_calls.append(sql)
6479+
6480+
return original_duckdb_execute(self, expression, **kwargs)
6481+
6482+
ctx = Context(paths=[tmp_path], config=config)
6483+
6484+
# the plan would fail if the execution order ever changes and before_all statements dont execute first
6485+
ctx.plan(auto_apply=True, no_prompts=True)
6486+
6487+
mocker.patch.object(DuckDBEngineAdapter, "execute", track_duckdb_execute)
6488+
6489+
# run with the patched execute
6490+
ctx.run("prod", start="2023-01-01", end="2023-01-02")
6491+
6492+
# validate explicitly that the first execute is for the before_all
6493+
assert "before_all_created_table" in execute_calls[0]
6494+
6495+
# and that the last is the sole after all that depends on the model
6496+
assert "after_all_created_table" in execute_calls[-1]
6497+
6498+
64306499
@time_machine.travel("2025-03-08 00:00:00 UTC")
64316500
def test_tz(init_and_plan_context):
64326501
context, _ = init_and_plan_context("examples/sushi")

0 commit comments

Comments
 (0)