Fix: Support of vars in on-run-start / on-run-end hooks in dbt projects (#4554)

Author: izeigerman

sqlmesh/core/environment.py

@@ -250,6 +250,39 @@ class EnvironmentStatements(PydanticModel):
     python_env: t.Dict[str, Executable]
     jinja_macros: t.Optional[JinjaMacroRegistry] = None
 
+    def render_before_all(
+        self,
+        dialect: str,
+        default_catalog: t.Optional[str] = None,
+        **render_kwargs: t.Any,
+    ) -> t.List[str]:
+        return self.render(RuntimeStage.BEFORE_ALL, dialect, default_catalog, **render_kwargs)
+
+    def render_after_all(
+        self,
+        dialect: str,
+        default_catalog: t.Optional[str] = None,
+        **render_kwargs: t.Any,
+    ) -> t.List[str]:
+        return self.render(RuntimeStage.AFTER_ALL, dialect, default_catalog, **render_kwargs)
+
+    def render(
+        self,
+        runtime_stage: RuntimeStage,
+        dialect: str,
+        default_catalog: t.Optional[str] = None,
+        **render_kwargs: t.Any,
+    ) -> t.List[str]:
+        return render_statements(
+            statements=getattr(self, runtime_stage.value),
+            dialect=dialect,
+            default_catalog=default_catalog,
+            python_env=self.python_env,
+            jinja_macros=self.jinja_macros,
+            runtime_stage=runtime_stage,
+            **render_kwargs,
+        )
+
 
 def execute_environment_statements(
     adapter: EngineAdapter,
@@ -266,18 +299,15 @@ def execute_environment_statements(
         rendered_expressions = [
             expr
             for statements in environment_statements
-            for expr in render_statements(
-                statements=getattr(statements, runtime_stage.value),
+            for expr in statements.render(
+                runtime_stage=runtime_stage,
                 dialect=adapter.dialect,
                 default_catalog=default_catalog,
-                python_env=statements.python_env,
-                jinja_macros=statements.jinja_macros,
                 snapshots=snapshots,
                 start=start,
                 end=end,
                 execution_time=execution_time,
                 environment_naming_info=environment_naming_info,
-                runtime_stage=runtime_stage,
                 engine_adapter=adapter,
             )
         ]

sqlmesh/dbt/loader.py

@@ -4,7 +4,6 @@
 import sys
 import typing as t
 import sqlmesh.core.dialect as d
-from sqlglot.optimizer.simplify import gen
 from pathlib import Path
 from sqlmesh.core import constants as c
 from sqlmesh.core.config import (
@@ -17,9 +16,9 @@
 from sqlmesh.core.loader import CacheBase, LoadedProject, Loader
 from sqlmesh.core.macros import MacroRegistry, macro
 from sqlmesh.core.model import Model, ModelCache
-from sqlmesh.core.model.common import make_python_env
 from sqlmesh.core.signal import signal
 from sqlmesh.dbt.basemodel import BMC, BaseModelConfig
+from sqlmesh.dbt.common import Dependencies
 from sqlmesh.dbt.context import DbtContext
 from sqlmesh.dbt.model import ModelConfig
 from sqlmesh.dbt.profile import Profile
@@ -259,22 +258,18 @@ def _load_environment_statements(self, macros: MacroRegistry) -> t.List[Environm
 
                 if statements := on_run_start + on_run_end:
                     jinja_references, used_variables = extract_macro_references_and_variables(
-                        *(gen(stmt) for stmt in statements)
+                        *statements
                     )
 
-                    jinja_registry = make_jinja_registry(
-                        context.jinja_macros, package_name, jinja_references
+                    statements_context = context.context_for_dependencies(
+                        Dependencies(
+                            variables=used_variables,
+                        )
                     )
-
-                    python_env = make_python_env(
-                        [s for stmt in statements for s in d.parse(stmt, default_dialect=dialect)],
-                        jinja_macro_references=jinja_references,
-                        module_path=self.config_path,
-                        macros=macros,
-                        variables=context.variables,
-                        used_variables=used_variables,
-                        path=self.config_path,
+                    jinja_registry = make_jinja_registry(
+                        statements_context.jinja_macros, package_name, jinja_references
                     )
+                    jinja_registry.add_globals(statements_context.jinja_globals)
 
                     hooks_by_package_name[package_name] = EnvironmentStatements(
                         before_all=[
@@ -285,7 +280,7 @@ def _load_environment_statements(self, macros: MacroRegistry) -> t.List[Environm
                             d.jinja_statement(stmt).sql(dialect=dialect)
                             for stmt in on_run_end or []
                         ],
-                        python_env=python_env,
+                        python_env={},
                         jinja_macros=jinja_registry,
                     )
                 # Project hooks should be executed first and then rest of the packages

tests/core/test_integration.py

@@ -4345,6 +4345,21 @@ def test_dbt_dialect_with_normalization_strategy(init_and_plan_context: t.Callab
     assert context.default_dialect == "duckdb,normalization_strategy=LOWERCASE"
 
 
+@time_machine.travel("2023-01-08 15:00:00 UTC")
+def test_dbt_before_all_with_var(init_and_plan_context: t.Callable):
+    _, plan = init_and_plan_context(
+        "tests/fixtures/dbt/sushi_test", config="test_config_with_normalization_strategy"
+    )
+    environment_statements = plan.to_evaluatable().environment_statements
+    assert environment_statements
+    rendered_statements = [e.render_before_all(dialect="duckdb") for e in environment_statements]
+    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"',
+    ]
+
+
 @pytest.mark.parametrize(
     "context_fixture",
     ["sushi_context", "sushi_dbt_context", "sushi_test_dbt_context", "sushi_no_default_catalog"],

tests/dbt/test_adapter.py

@@ -288,6 +288,7 @@ def test_on_run_start_end(copy_to_temp_path):
     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;",
@@ -317,6 +318,7 @@ def test_on_run_start_end(copy_to_temp_path):
     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

tests/fixtures/dbt/sushi_test/dbt_project.yml

@@ -62,6 +62,7 @@ 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") }}'
 on-run-end:
   - '{{ create_tables(schemas) }}'
   - 'DROP TABLE to_be_executed_last;'