Fix: Dependency handling when converting on-run-start / on-run-end hooks in dbt projects (#4567)

izeigerman · web-flow · commit 2785fc98615c · 2025-05-28T10:00:40.000-07:00
diff --git a/sqlmesh/dbt/loader.py b/sqlmesh/dbt/loader.py
@@ -28,7 +28,6 @@
 from sqlmesh.utils.errors import ConfigError
 from sqlmesh.utils.jinja import (
     JinjaMacroRegistry,
-    extract_macro_references_and_variables,
     make_jinja_registry,
 )
 
@@ -240,11 +239,11 @@ def _load_requirements(self) -> t.Tuple[t.Dict[str, str], t.Set[str]]:
     def _load_environment_statements(self, macros: MacroRegistry) -> t.List[EnvironmentStatements]:
         """Loads dbt's on_run_start, on_run_end hooks into sqlmesh's before_all, after_all statements respectively."""
 
-        environment_statements: t.List[EnvironmentStatements] = []
+        hooks_by_package_name: t.Dict[str, EnvironmentStatements] = {}
+        project_names: t.Set[str] = set()
         dialect = self.config.dialect
         for project in self._load_projects():
             context = project.context
-            hooks_by_package_name: t.Dict[str, EnvironmentStatements] = {}
             for package_name, package in project.packages.items():
                 context.set_and_render_variables(package.variables, package_name)
                 on_run_start: t.List[str] = [
@@ -256,18 +255,14 @@ def _load_environment_statements(self, macros: MacroRegistry) -> t.List[Environm
                     for on_run_hook in sorted(package.on_run_end.values(), key=lambda h: h.index)
                 ]
 
-                if statements := on_run_start + on_run_end:
-                    jinja_references, used_variables = extract_macro_references_and_variables(
-                        *statements
-                    )
+                if on_run_start or on_run_end:
+                    dependencies = Dependencies()
+                    for hook in [*package.on_run_start.values(), *package.on_run_end.values()]:
+                        dependencies = dependencies.union(hook.dependencies)
 
-                    statements_context = context.context_for_dependencies(
-                        Dependencies(
-                            variables=used_variables,
-                        )
-                    )
+                    statements_context = context.context_for_dependencies(dependencies)
                     jinja_registry = make_jinja_registry(
-                        statements_context.jinja_macros, package_name, jinja_references
+                        statements_context.jinja_macros, package_name, set(dependencies.macros)
                     )
                     jinja_registry.add_globals(statements_context.jinja_globals)
 
@@ -283,15 +278,15 @@ def _load_environment_statements(self, macros: MacroRegistry) -> t.List[Environm
                         python_env={},
                         jinja_macros=jinja_registry,
                     )
-                # Project hooks should be executed first and then rest of the packages
-                environment_statements = [
-                    statements
-                    for _, statements in sorted(
-                        hooks_by_package_name.items(),
-                        key=lambda item: 0 if item[0] == context.project_name else 1,
-                    )
-                ]
-        return environment_statements
+                    project_names.add(package_name)
+
+        return [
+            statements
+            for _, statements in sorted(
+                hooks_by_package_name.items(),
+                key=lambda item: 0 if item[0] in project_names else 1,
+            )
+        ]
 
     def _compute_yaml_max_mtime_per_subfolder(self, root: Path) -> t.Dict[Path, float]:
         if not root.is_dir():
diff --git a/sqlmesh/dbt/manifest.py b/sqlmesh/dbt/manifest.py
@@ -292,13 +292,24 @@ def _load_on_run_start_end(self) -> None:
                 sql = node.raw_code if DBT_VERSION >= (1, 3) else node.raw_sql  # type: ignore
                 node_name = node.name
                 node_path = Path(node.original_file_path)
+
+                dependencies = Dependencies(
+                    macros=_macro_references(self._manifest, node),
+                    refs=_refs(node),
+                    sources=_sources(node),
+                )
+                dependencies = dependencies.union(self._extra_dependencies(sql, node.package_name))
+                dependencies = dependencies.union(
+                    self._flatten_dependencies_from_macros(dependencies.macros, node.package_name)
+                )
+
                 if "on-run-start" in node.tags:
                     self._on_run_start_per_package[node.package_name][node_name] = HookConfig(
-                        sql=sql, index=node.index or 0, path=node_path
+                        sql=sql, index=node.index or 0, path=node_path, dependencies=dependencies
                     )
                 else:
                     self._on_run_end_per_package[node.package_name][node_name] = HookConfig(
-                        sql=sql, index=node.index or 0, path=node_path
+                        sql=sql, index=node.index or 0, path=node_path, dependencies=dependencies
                     )
 
     @property
diff --git a/sqlmesh/dbt/package.py b/sqlmesh/dbt/package.py
@@ -34,6 +34,7 @@ class HookConfig(PydanticModel):
     sql: str
     index: int
     path: Path
+    dependencies: Dependencies
 
 
 class Package(PydanticModel):
diff --git a/tests/core/test_integration.py b/tests/core/test_integration.py
@@ -4346,7 +4346,7 @@ def test_dbt_dialect_with_normalization_strategy(init_and_plan_context: t.Callab
 
 
 @time_machine.travel("2023-01-08 15:00:00 UTC")
-def test_dbt_before_all_with_var(init_and_plan_context: t.Callable):
+def test_dbt_before_all_with_var_ref_source(init_and_plan_context: t.Callable):
     _, plan = init_and_plan_context(
         "tests/fixtures/dbt/sushi_test", config="test_config_with_normalization_strategy"
     )
@@ -4356,7 +4356,7 @@ def test_dbt_before_all_with_var(init_and_plan_context: t.Callable):
     assert rendered_statements[0] == [
         "CREATE TABLE IF NOT EXISTS analytic_stats (physical_table TEXT, evaluation_time TEXT)",
         "CREATE TABLE IF NOT EXISTS to_be_executed_last (col TEXT)",
-        'SELECT 1 AS "1"',
+        "SELECT 1 AS var, 'items' AS src, 'waiters' AS ref",
     ]
 
 
diff --git a/tests/dbt/test_adapter.py b/tests/dbt/test_adapter.py
@@ -13,9 +13,6 @@
 
 from sqlmesh import Context
 from sqlmesh.core.dialect import schema_
-from sqlmesh.core.environment import EnvironmentNamingInfo
-from sqlmesh.core.macros import RuntimeStage
-from sqlmesh.core.renderer import render_statements
 from sqlmesh.core.snapshot import SnapshotId
 from sqlmesh.dbt.adapter import ParsetimeAdapter
 from sqlmesh.dbt.project import Project
@@ -275,114 +272,6 @@ def test_quote_as_configured():
     adapter.quote_as_configured("foo", "database") == "foo"
 
 
-def test_on_run_start_end(copy_to_temp_path):
-    project_root = "tests/fixtures/dbt/sushi_test"
-    sushi_context = Context(paths=copy_to_temp_path(project_root))
-    assert len(sushi_context._environment_statements) == 2
-
-    # Root project's on run start / on run end should be first by checking the macros
-    root_environment_statements = sushi_context._environment_statements[0]
-    assert "create_tables" in root_environment_statements.jinja_macros.root_macros
-
-    # Validate order of execution to be correct
-    assert root_environment_statements.before_all == [
-        "JINJA_STATEMENT_BEGIN;\nCREATE TABLE IF NOT EXISTS analytic_stats (physical_table VARCHAR, evaluation_time VARCHAR);\nJINJA_END;",
-        "JINJA_STATEMENT_BEGIN;\nCREATE TABLE IF NOT EXISTS to_be_executed_last (col VARCHAR);\nJINJA_END;",
-        'JINJA_STATEMENT_BEGIN;\nSELECT {{ var("yet_another_var") }}\nJINJA_END;',
-    ]
-    assert root_environment_statements.after_all == [
-        "JINJA_STATEMENT_BEGIN;\n{{ create_tables(schemas) }}\nJINJA_END;",
-        "JINJA_STATEMENT_BEGIN;\nDROP TABLE to_be_executed_last;\nJINJA_END;",
-    ]
-
-    assert root_environment_statements.jinja_macros.root_package_name == "sushi"
-
-    rendered_before_all = render_statements(
-        root_environment_statements.before_all,
-        dialect=sushi_context.default_dialect,
-        python_env=root_environment_statements.python_env,
-        jinja_macros=root_environment_statements.jinja_macros,
-        runtime_stage=RuntimeStage.BEFORE_ALL,
-    )
-
-    rendered_after_all = render_statements(
-        root_environment_statements.after_all,
-        dialect=sushi_context.default_dialect,
-        python_env=root_environment_statements.python_env,
-        jinja_macros=root_environment_statements.jinja_macros,
-        snapshots=sushi_context.snapshots,
-        runtime_stage=RuntimeStage.AFTER_ALL,
-        environment_naming_info=EnvironmentNamingInfo(name="dev"),
-    )
-
-    assert rendered_before_all == [
-        "CREATE TABLE IF NOT EXISTS analytic_stats (physical_table TEXT, evaluation_time TEXT)",
-        "CREATE TABLE IF NOT EXISTS to_be_executed_last (col TEXT)",
-        'SELECT 1 AS "1"',
-    ]
-
-    # The jinja macro should have resolved the schemas for this environment and generated corresponding statements
-    assert sorted(rendered_after_all) == sorted(
-        [
-            "CREATE OR REPLACE TABLE schema_table_snapshots__dev AS SELECT 'snapshots__dev' AS schema",
-            "CREATE OR REPLACE TABLE schema_table_sushi__dev AS SELECT 'sushi__dev' AS schema",
-            "DROP TABLE to_be_executed_last",
-        ]
-    )
-
-    # Nested dbt_packages on run start / on run end
-    packaged_environment_statements = sushi_context._environment_statements[1]
-
-    # Validate order of execution to be correct
-    assert packaged_environment_statements.before_all == [
-        "JINJA_STATEMENT_BEGIN;\nCREATE TABLE IF NOT EXISTS to_be_executed_first (col VARCHAR);\nJINJA_END;",
-        "JINJA_STATEMENT_BEGIN;\nCREATE TABLE IF NOT EXISTS analytic_stats_packaged_project (physical_table VARCHAR, evaluation_time VARCHAR);\nJINJA_END;",
-    ]
-    assert packaged_environment_statements.after_all == [
-        "JINJA_STATEMENT_BEGIN;\nDROP TABLE to_be_executed_first\nJINJA_END;",
-        "JINJA_STATEMENT_BEGIN;\n{{ packaged_tables(schemas) }}\nJINJA_END;",
-    ]
-
-    assert "packaged_tables" in packaged_environment_statements.jinja_macros.root_macros
-    assert packaged_environment_statements.jinja_macros.root_package_name == "sushi"
-
-    rendered_before_all = render_statements(
-        packaged_environment_statements.before_all,
-        dialect=sushi_context.default_dialect,
-        python_env=packaged_environment_statements.python_env,
-        jinja_macros=packaged_environment_statements.jinja_macros,
-        runtime_stage=RuntimeStage.BEFORE_ALL,
-    )
-
-    rendered_after_all = render_statements(
-        packaged_environment_statements.after_all,
-        dialect=sushi_context.default_dialect,
-        python_env=packaged_environment_statements.python_env,
-        jinja_macros=packaged_environment_statements.jinja_macros,
-        snapshots=sushi_context.snapshots,
-        runtime_stage=RuntimeStage.AFTER_ALL,
-        environment_naming_info=EnvironmentNamingInfo(name="dev"),
-    )
-
-    # Validate order of execution to match dbt's
-    assert rendered_before_all == [
-        "CREATE TABLE IF NOT EXISTS to_be_executed_first (col TEXT)",
-        "CREATE TABLE IF NOT EXISTS analytic_stats_packaged_project (physical_table TEXT, evaluation_time TEXT)",
-    ]
-
-    # This on run end statement should be executed first
-    assert rendered_after_all[0] == "DROP TABLE to_be_executed_first"
-
-    # The table names is an indication of the rendering of the dbt_packages statements
-    assert sorted(rendered_after_all) == sorted(
-        [
-            "DROP TABLE to_be_executed_first",
-            "CREATE OR REPLACE TABLE schema_table_snapshots__dev_nested_package AS SELECT 'snapshots__dev' AS schema",
-            "CREATE OR REPLACE TABLE schema_table_sushi__dev_nested_package AS SELECT 'sushi__dev' AS schema",
-        ]
-    )
-
-
 def test_adapter_get_relation_normalization(
     sushi_test_project: Project, runtime_renderer: t.Callable
 ):
diff --git a/tests/dbt/test_transformation.py b/tests/dbt/test_transformation.py
@@ -13,6 +13,9 @@
 from pytest_mock.plugin import MockerFixture
 from sqlglot import exp, parse_one
 from sqlmesh.core import dialect as d
+from sqlmesh.core.environment import EnvironmentNamingInfo
+from sqlmesh.core.macros import RuntimeStage
+from sqlmesh.core.renderer import render_statements
 from sqlmesh.core.audit import StandaloneAudit
 from sqlmesh.core.context import Context
 from sqlmesh.core.console import get_console
@@ -1551,3 +1554,112 @@ def test_grain():
 
     model.grain = "id_a"
     assert model.to_sqlmesh(context).grains == [exp.to_column("id_a")]
+
+
+def test_on_run_start_end(copy_to_temp_path):
+    project_root = "tests/fixtures/dbt/sushi_test"
+    sushi_context = Context(paths=copy_to_temp_path(project_root))
+    assert len(sushi_context._environment_statements) == 2
+
+    # Root project's on run start / on run end should be first by checking the macros
+    root_environment_statements = sushi_context._environment_statements[0]
+    assert "create_tables" in root_environment_statements.jinja_macros.root_macros
+
+    # Validate order of execution to be correct
+    assert root_environment_statements.before_all == [
+        "JINJA_STATEMENT_BEGIN;\nCREATE TABLE IF NOT EXISTS analytic_stats (physical_table VARCHAR, evaluation_time VARCHAR);\nJINJA_END;",
+        "JINJA_STATEMENT_BEGIN;\nCREATE TABLE IF NOT EXISTS to_be_executed_last (col VARCHAR);\nJINJA_END;",
+        """JINJA_STATEMENT_BEGIN;\nSELECT {{ var("yet_another_var") }} AS var, '{{ source("raw", "items").identifier }}' AS src, '{{ ref("waiters").identifier }}' AS ref;\nJINJA_END;""",
+        "JINJA_STATEMENT_BEGIN;\n{{ log_value('on-run-start') }}\nJINJA_END;",
+    ]
+    assert root_environment_statements.after_all == [
+        "JINJA_STATEMENT_BEGIN;\n{{ create_tables(schemas) }}\nJINJA_END;",
+        "JINJA_STATEMENT_BEGIN;\nDROP TABLE to_be_executed_last;\nJINJA_END;",
+    ]
+
+    assert root_environment_statements.jinja_macros.root_package_name == "sushi"
+
+    rendered_before_all = render_statements(
+        root_environment_statements.before_all,
+        dialect=sushi_context.default_dialect,
+        python_env=root_environment_statements.python_env,
+        jinja_macros=root_environment_statements.jinja_macros,
+        runtime_stage=RuntimeStage.BEFORE_ALL,
+    )
+
+    rendered_after_all = render_statements(
+        root_environment_statements.after_all,
+        dialect=sushi_context.default_dialect,
+        python_env=root_environment_statements.python_env,
+        jinja_macros=root_environment_statements.jinja_macros,
+        snapshots=sushi_context.snapshots,
+        runtime_stage=RuntimeStage.AFTER_ALL,
+        environment_naming_info=EnvironmentNamingInfo(name="dev"),
+    )
+
+    assert rendered_before_all == [
+        "CREATE TABLE IF NOT EXISTS analytic_stats (physical_table TEXT, evaluation_time TEXT)",
+        "CREATE TABLE IF NOT EXISTS to_be_executed_last (col TEXT)",
+        "SELECT 1 AS var, 'items' AS src, 'waiters' AS ref",
+    ]
+
+    # The jinja macro should have resolved the schemas for this environment and generated corresponding statements
+    assert sorted(rendered_after_all) == sorted(
+        [
+            "CREATE OR REPLACE TABLE schema_table_snapshots__dev AS SELECT 'snapshots__dev' AS schema",
+            "CREATE OR REPLACE TABLE schema_table_sushi__dev AS SELECT 'sushi__dev' AS schema",
+            "DROP TABLE to_be_executed_last",
+        ]
+    )
+
+    # Nested dbt_packages on run start / on run end
+    packaged_environment_statements = sushi_context._environment_statements[1]
+
+    # Validate order of execution to be correct
+    assert packaged_environment_statements.before_all == [
+        "JINJA_STATEMENT_BEGIN;\nCREATE TABLE IF NOT EXISTS to_be_executed_first (col VARCHAR);\nJINJA_END;",
+        "JINJA_STATEMENT_BEGIN;\nCREATE TABLE IF NOT EXISTS analytic_stats_packaged_project (physical_table VARCHAR, evaluation_time VARCHAR);\nJINJA_END;",
+    ]
+    assert packaged_environment_statements.after_all == [
+        "JINJA_STATEMENT_BEGIN;\nDROP TABLE to_be_executed_first\nJINJA_END;",
+        "JINJA_STATEMENT_BEGIN;\n{{ packaged_tables(schemas) }}\nJINJA_END;",
+    ]
+
+    assert "packaged_tables" in packaged_environment_statements.jinja_macros.root_macros
+    assert packaged_environment_statements.jinja_macros.root_package_name == "sushi"
+
+    rendered_before_all = render_statements(
+        packaged_environment_statements.before_all,
+        dialect=sushi_context.default_dialect,
+        python_env=packaged_environment_statements.python_env,
+        jinja_macros=packaged_environment_statements.jinja_macros,
+        runtime_stage=RuntimeStage.BEFORE_ALL,
+    )
+
+    rendered_after_all = render_statements(
+        packaged_environment_statements.after_all,
+        dialect=sushi_context.default_dialect,
+        python_env=packaged_environment_statements.python_env,
+        jinja_macros=packaged_environment_statements.jinja_macros,
+        snapshots=sushi_context.snapshots,
+        runtime_stage=RuntimeStage.AFTER_ALL,
+        environment_naming_info=EnvironmentNamingInfo(name="dev"),
+    )
+
+    # Validate order of execution to match dbt's
+    assert rendered_before_all == [
+        "CREATE TABLE IF NOT EXISTS to_be_executed_first (col TEXT)",
+        "CREATE TABLE IF NOT EXISTS analytic_stats_packaged_project (physical_table TEXT, evaluation_time TEXT)",
+    ]
+
+    # This on run end statement should be executed first
+    assert rendered_after_all[0] == "DROP TABLE to_be_executed_first"
+
+    # The table names is an indication of the rendering of the dbt_packages statements
+    assert sorted(rendered_after_all) == sorted(
+        [
+            "DROP TABLE to_be_executed_first",
+            "CREATE OR REPLACE TABLE schema_table_snapshots__dev_nested_package AS SELECT 'snapshots__dev' AS schema",
+            "CREATE OR REPLACE TABLE schema_table_sushi__dev_nested_package AS SELECT 'sushi__dev' AS schema",
+        ]
+    )
diff --git a/tests/fixtures/dbt/sushi_test/dbt_project.yml b/tests/fixtures/dbt/sushi_test/dbt_project.yml
@@ -62,7 +62,8 @@ vars:
 on-run-start:
   - 'CREATE TABLE IF NOT EXISTS analytic_stats (physical_table VARCHAR, evaluation_time VARCHAR);'
   - 'CREATE TABLE IF NOT EXISTS to_be_executed_last (col VARCHAR);'
-  - 'SELECT {{ var("yet_another_var") }}'
+  - SELECT {{ var("yet_another_var") }} AS var, '{{ source("raw", "items").identifier }}' AS src, '{{ ref("waiters").identifier }}' AS ref;
+  - "{{ log_value('on-run-start') }}"
 on-run-end:
   - '{{ create_tables(schemas) }}'
   - 'DROP TABLE to_be_executed_last;'

Original file line number	Diff line number	Diff line change
`@@ -4346,7 +4346,7 @@ def test_dbt_dialect_with_normalization_strategy(init_and_plan_context: t.Callab`
`4346`	`4346`
`4347`	`4347`
`4348`	`4348`	`@time_machine.travel("2023-01-08 15:00:00 UTC")`
`4349`		`-def test_dbt_before_all_with_var(init_and_plan_context: t.Callable):`
	`4349`	`+def test_dbt_before_all_with_var_ref_source(init_and_plan_context: t.Callable):`
`4350`	`4350`	`_, plan = init_and_plan_context(`
`4351`	`4351`	`"tests/fixtures/dbt/sushi_test", config="test_config_with_normalization_strategy"`
`4352`	`4352`	`)`
`@@ -4356,7 +4356,7 @@ def test_dbt_before_all_with_var(init_and_plan_context: t.Callable):`
`4356`	`4356`	`assert rendered_statements[0] == [`
`4357`	`4357`	`"CREATE TABLE IF NOT EXISTS analytic_stats (physical_table TEXT, evaluation_time TEXT)",`
`4358`	`4358`	`"CREATE TABLE IF NOT EXISTS to_be_executed_last (col TEXT)",`
`4359`		`- 'SELECT 1 AS "1"',`
	`4359`	`+ "SELECT 1 AS var, 'items' AS src, 'waiters' AS ref",`
`4360`	`4360`	`]`
`4361`	`4361`
`4362`	`4362`