Fix: Include on run start / on run end hooks of dbt packages

Author: themisvaltinos

sqlmesh/core/context.py

@@ -610,8 +610,7 @@ def load(self, update_schemas: bool = True) -> GenericContext[C]:
             self._standalone_audits.update(project.standalone_audits)
             self._requirements.update(project.requirements)
             self._excluded_requirements.update(project.excluded_requirements)
-            if project.environment_statements:
-                self._environment_statements.append(project.environment_statements)
+            self._environment_statements.extend(project.environment_statements)
 
             config = loader.config
             self._linters[config.project] = Linter.from_rules(

sqlmesh/core/loader.py

@@ -61,7 +61,7 @@ class LoadedProject:
     metrics: UniqueKeyDict[str, Metric]
     requirements: t.Dict[str, str]
     excluded_requirements: t.Set[str]
-    environment_statements: t.Optional[EnvironmentStatements]
+    environment_statements: t.List[EnvironmentStatements]
     user_rules: RuleSet
 
 
@@ -187,9 +187,9 @@ def _load_audits(
     ) -> UniqueKeyDict[str, Audit]:
         """Loads all audits."""
 
-    def _load_environment_statements(self, macros: MacroRegistry) -> EnvironmentStatements | None:
+    def _load_environment_statements(self, macros: MacroRegistry) -> t.List[EnvironmentStatements]:
         """Loads environment statements."""
-        return None
+        return []
 
     def load_materializations(self) -> None:
         """Loads custom materializations."""
@@ -651,7 +651,7 @@ def _load_metrics(self) -> UniqueKeyDict[str, MetricMeta]:
 
         return metrics
 
-    def _load_environment_statements(self, macros: MacroRegistry) -> EnvironmentStatements | None:
+    def _load_environment_statements(self, macros: MacroRegistry) -> t.List[EnvironmentStatements]:
         """Loads environment statements."""
 
         if self.config.before_all or self.config.after_all:
@@ -673,8 +673,8 @@ def _load_environment_statements(self, macros: MacroRegistry) -> EnvironmentStat
                 path=self.config_path,
             )
 
-            return EnvironmentStatements(**statements, python_env=python_env)
-        return None
+            return [EnvironmentStatements(**statements, python_env=python_env)]
+        return []
 
     def _load_linting_rules(self) -> RuleSet:
         user_rules: UniqueKeyDict[str, type[Rule]] = UniqueKeyDict("rules")

sqlmesh/dbt/loader.py

@@ -29,7 +29,6 @@
 from sqlmesh.utils.errors import ConfigError
 from sqlmesh.utils.jinja import (
     JinjaMacroRegistry,
-    MacroInfo,
     extract_macro_references_and_variables,
 )
 
@@ -238,59 +237,54 @@ def _load_requirements(self) -> t.Tuple[t.Dict[str, str], t.Set[str]]:
 
         return requirements, excluded_requirements
 
-    def _load_environment_statements(self, macros: MacroRegistry) -> EnvironmentStatements | None:
+    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."""
 
-        on_run_start: t.List[str] = []
-        on_run_end: t.List[str] = []
-        jinja_root_macros: t.Dict[str, MacroInfo] = {}
-        variables: t.Dict[str, t.Any] = self._get_variables()
+        environment_statements: t.List[EnvironmentStatements] = []
         dialect = self.config.dialect
         for project in self._load_projects():
-            context = project.context.copy()
-            if manifest := context._manifest:
-                on_run_start.extend(manifest._on_run_start or [])
-                on_run_end.extend(manifest._on_run_end or [])
-
-            if root_package := context.jinja_macros.root_package_name:
-                if root_macros := context.jinja_macros.packages.get(root_package):
-                    jinja_root_macros |= root_macros
-                context.set_and_render_variables(context.variables, root_package)
-                variables |= context.variables
-
-        if statements := on_run_start + on_run_end:
-            jinja_macro_references, used_variables = extract_macro_references_and_variables(
-                *(gen(stmt) for stmt in statements)
-            )
-            jinja_macros = context.jinja_macros
-            jinja_macros.root_macros = jinja_root_macros
-            jinja_macros = (
-                jinja_macros.trim(jinja_macro_references)
-                if not jinja_macros.trimmed
-                else jinja_macros
-            )
-
-            python_env = make_python_env(
-                [s for stmt in statements for s in d.parse(stmt, default_dialect=dialect)],
-                jinja_macro_references=jinja_macro_references,
-                module_path=self.config_path,
-                macros=macros,
-                variables=variables,
-                used_variables=used_variables,
-                path=self.config_path,
-            )
-
-            return EnvironmentStatements(
-                before_all=[
-                    d.jinja_statement(stmt).sql(dialect=dialect) for stmt in on_run_start or []
-                ],
-                after_all=[
-                    d.jinja_statement(stmt).sql(dialect=dialect) for stmt in on_run_end or []
-                ],
-                python_env=python_env,
-                jinja_macros=jinja_macros,
-            )
-        return None
+            context = project.context
+            for package_name, package in project.packages.items():
+                context.set_and_render_variables(package.variables, package_name)
+                on_run_start: t.List[str] = []
+                on_run_end: t.List[str] = []
+                for hook in package.on_run_start.values():
+                    on_run_start.append(hook.sql)
+                for hook in package.on_run_end.values():
+                    on_run_end.append(hook.sql)
+
+                if statements := on_run_start + on_run_end:
+                    jinja_references, used_variables = extract_macro_references_and_variables(
+                        *(gen(stmt) for stmt in statements)
+                    )
+                    jinja_registry = context.jinja_macros.copy()
+                    jinja_registry.root_macros = jinja_registry.packages.get(package_name) or {}
+                    jinja_registry = jinja_registry.trim(jinja_references)
+                    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,
+                    )
+
+                    environment_statements.append(
+                        EnvironmentStatements(
+                            before_all=[
+                                d.jinja_statement(stmt).sql(dialect=dialect)
+                                for stmt in on_run_start or []
+                            ],
+                            after_all=[
+                                d.jinja_statement(stmt).sql(dialect=dialect)
+                                for stmt in on_run_end or []
+                            ],
+                            python_env=python_env,
+                            jinja_macros=jinja_registry,
+                        )
+                    )
+        return environment_statements
 
     def _compute_yaml_max_mtime_per_subfolder(self, root: Path) -> t.Dict[Path, float]:
         if not root.is_dir():

sqlmesh/dbt/manifest.py

@@ -27,7 +27,7 @@
 from sqlmesh.dbt.basemodel import Dependencies
 from sqlmesh.dbt.builtin import BUILTIN_FILTERS, BUILTIN_GLOBALS, OVERRIDDEN_MACROS
 from sqlmesh.dbt.model import ModelConfig
-from sqlmesh.dbt.package import MacroConfig
+from sqlmesh.dbt.package import HookConfig, MacroConfig
 from sqlmesh.dbt.seed import SeedConfig
 from sqlmesh.dbt.source import SourceConfig
 from sqlmesh.dbt.target import TargetConfig
@@ -54,6 +54,7 @@
 SeedConfigs = t.Dict[str, SeedConfig]
 SourceConfigs = t.Dict[str, SourceConfig]
 MacroConfigs = t.Dict[str, MacroConfig]
+HookConfigs = t.Dict[str, HookConfig]
 
 
 IGNORED_PACKAGES = {"elementary"}
@@ -94,8 +95,8 @@ def __init__(
             self.project_path / c.CACHE, "jinja_calls"
         )
 
-        self._on_run_start: t.Optional[t.List[str]] = None
-        self._on_run_end: t.Optional[t.List[str]] = None
+        self._on_run_start_per_package: t.Dict[str, HookConfigs] = defaultdict(dict)
+        self._on_run_end_per_package: t.Dict[str, HookConfigs] = defaultdict(dict)
 
     def tests(self, package_name: t.Optional[str] = None) -> TestConfigs:
         self._load_all()
@@ -117,6 +118,14 @@ def macros(self, package_name: t.Optional[str] = None) -> MacroConfigs:
         self._load_all()
         return self._macros_per_package[package_name or self._project_name]
 
+    def on_run_start(self, package_name: t.Optional[str] = None) -> HookConfigs:
+        self._load_all()
+        return self._on_run_start_per_package[package_name or self._project_name]
+
+    def on_run_end(self, package_name: t.Optional[str] = None) -> HookConfigs:
+        self._load_all()
+        return self._on_run_end_per_package[package_name or self._project_name]
+
     @property
     def all_macros(self) -> t.Dict[str, t.Dict[str, MacroInfo]]:
         self._load_all()
@@ -136,6 +145,7 @@ def _load_all(self) -> None:
         self._load_sources()
         self._load_tests()
         self._load_models_and_seeds()
+        self._load_on_run_start_end()
         self._is_loaded = True
 
         self._call_cache.put("", value={k: v for k, (v, used) in self._calls.items() if used})
@@ -274,6 +284,23 @@ def _load_models_and_seeds(self) -> None:
                     **node_config,
                 )
 
+    def _load_on_run_start_end(self) -> None:
+        for node in self._manifest.nodes.values():
+            if node.resource_type == "operation" and (
+                set(node.tags) & {"on-run-start", "on-run-end"}
+            ):
+                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)
+                if "on-run-start" in node.tags:
+                    self._on_run_start_per_package[node.package_name][node_name] = HookConfig(
+                        sql=sql, path=node_path
+                    )
+                else:
+                    self._on_run_end_per_package[node.package_name][node_name] = HookConfig(
+                        sql=sql, path=node_path
+                    )
+
     @property
     def _manifest(self) -> Manifest:
         if not self.__manifest:
@@ -315,11 +342,6 @@ def _load_manifest(self) -> Manifest:
 
         runtime_config = RuntimeConfig.from_parts(project, profile, args)
 
-        if runtime_config.on_run_start:
-            self._on_run_start = runtime_config.on_run_start
-        if runtime_config.on_run_end:
-            self._on_run_end = runtime_config.on_run_end
-
         self._project_name = project.project_name
 
         if DBT_VERSION >= (1, 8):

sqlmesh/dbt/package.py

@@ -28,6 +28,13 @@ class MacroConfig(PydanticModel):
     path: Path
 
 
+class HookConfig(PydanticModel):
+    """Class to contain on run start / on run end hooks."""
+
+    sql: str
+    path: Path
+
+
 class Package(PydanticModel):
     """Class to contain package configuration"""
 
@@ -38,6 +45,8 @@ class Package(PydanticModel):
     models: t.Dict[str, ModelConfig]
     variables: t.Dict[str, t.Any]
     macros: t.Dict[str, MacroConfig]
+    on_run_start: t.Dict[str, HookConfig]
+    on_run_end: t.Dict[str, HookConfig]
     files: t.Set[Path]
 
     @property
@@ -83,6 +92,8 @@ def load(self, package_root: Path) -> Package:
         models = _fix_paths(self._context.manifest.models(package_name), package_root)
         seeds = _fix_paths(self._context.manifest.seeds(package_name), package_root)
         macros = _fix_paths(self._context.manifest.macros(package_name), package_root)
+        on_run_start = _fix_paths(self._context.manifest.on_run_start(package_name), package_root)
+        on_run_end = _fix_paths(self._context.manifest.on_run_end(package_name), package_root)
         sources = self._context.manifest.sources(package_name)
 
         config_paths = {
@@ -102,10 +113,12 @@ def load(self, package_root: Path) -> Package:
             variables=package_variables,
             macros=macros,
             files=config_paths,
+            on_run_start=on_run_start,
+            on_run_end=on_run_end,
         )
 
 
-T = t.TypeVar("T", TestConfig, ModelConfig, MacroConfig, SeedConfig)
+T = t.TypeVar("T", TestConfig, ModelConfig, MacroConfig, SeedConfig, HookConfig)
 
 
 def _fix_paths(configs: t.Dict[str, T], package_root: Path) -> t.Dict[str, T]:

tests/dbt/test_adapter.py

@@ -278,30 +278,33 @@ def test_quote_as_configured():
 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) == 1
-    environment_statements = sushi_context._environment_statements[0]
+    assert len(sushi_context._environment_statements) == 2
 
-    assert environment_statements.before_all == [
+    # Root project on run start / on run end
+    root_environment_statements = sushi_context._environment_statements[0]
+    assert root_environment_statements.before_all == [
         "JINJA_STATEMENT_BEGIN;\nCREATE TABLE IF NOT EXISTS analytic_stats (physical_table VARCHAR, evaluation_time VARCHAR);\nJINJA_END;"
     ]
-    assert environment_statements.after_all == [
+    assert root_environment_statements.after_all == [
         "JINJA_STATEMENT_BEGIN;\n{{ create_tables(schemas) }}\nJINJA_END;"
     ]
-    assert "create_tables" in environment_statements.jinja_macros.root_macros
+
+    assert "create_tables" in root_environment_statements.jinja_macros.root_macros
+    assert root_environment_statements.jinja_macros.root_package_name == "sushi"
 
     rendered_before_all = render_statements(
-        environment_statements.before_all,
+        root_environment_statements.before_all,
         dialect=sushi_context.default_dialect,
-        python_env=environment_statements.python_env,
-        jinja_macros=environment_statements.jinja_macros,
+        python_env=root_environment_statements.python_env,
+        jinja_macros=root_environment_statements.jinja_macros,
         runtime_stage=RuntimeStage.BEFORE_ALL,
     )
 
     rendered_after_all = render_statements(
-        environment_statements.after_all,
+        root_environment_statements.after_all,
         dialect=sushi_context.default_dialect,
-        python_env=environment_statements.python_env,
-        jinja_macros=environment_statements.jinja_macros,
+        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"),
@@ -318,3 +321,46 @@ def test_on_run_start_end(copy_to_temp_path):
             "CREATE OR REPLACE TABLE schema_table_sushi__dev AS SELECT 'sushi__dev' AS schema",
         ]
     )
+
+    # Nested dbt_packages on run start / on run end
+    packaged_environment_statements = sushi_context._environment_statements[1]
+
+    assert packaged_environment_statements.before_all == [
+        "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;\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"),
+    )
+
+    assert rendered_before_all == [
+        "CREATE TABLE IF NOT EXISTS analytic_stats_packaged_project (physical_table TEXT, evaluation_time TEXT)"
+    ]
+
+    # The table names is an indication of the rendering of the dbt_packages statements
+    assert sorted(rendered_after_all) == sorted(
+        [
+            "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",
+        ]
+    )