Fix: Support on run start / on run end hooks of dbt packages (#4222)

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,8 +29,8 @@
 from sqlmesh.utils.errors import ConfigError
 from sqlmesh.utils.jinja import (
     JinjaMacroRegistry,
-    MacroInfo,
     extract_macro_references_and_variables,
+    make_jinja_registry,
 )
 
 if sys.version_info >= (3, 12):
@@ -238,59 +238,65 @@ 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,
-            )
+            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] = [
+                    on_run_hook.sql
+                    for on_run_hook in sorted(package.on_run_start.values(), key=lambda h: h.index)
+                ]
+                on_run_end: t.List[str] = [
+                    on_run_hook.sql
+                    for on_run_hook in sorted(package.on_run_end.values(), key=lambda h: h.index)
+                ]
 
-            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
+                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 = make_jinja_registry(
+                        context.jinja_macros, package_name, 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,
+                    )
+
+                    hooks_by_package_name[package_name] = 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,
+                    )
+                # 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
 
     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, index=node.index or 0, path=node_path
+                    )
+                else:
+                    self._on_run_end_per_package[node.package_name][node_name] = HookConfig(
+                        sql=sql, index=node.index or 0, 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,14 @@ class MacroConfig(PydanticModel):
     path: Path
 
 
+class HookConfig(PydanticModel):
+    """Class to contain on run start / on run end hooks."""
+
+    sql: str
+    index: int
+    path: Path
+
+
 class Package(PydanticModel):
     """Class to contain package configuration"""
 
@@ -38,6 +46,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 +93,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 +114,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]:

sqlmesh/utils/jinja.py

@@ -608,3 +608,31 @@ def create_builtin_globals(
         c.GATEWAY: lambda: variables.get(c.GATEWAY, None),
         **global_vars,
     }
+
+
+def make_jinja_registry(
+    jinja_macros: JinjaMacroRegistry, package_name: str, jinja_references: t.Set[MacroReference]
+) -> JinjaMacroRegistry:
+    """
+    Creates a Jinja macro registry for a specific package.
+
+    This function takes an existing Jinja macro registry and returns a new
+    registry that includes only the macros associated with the specified
+    package and trims the registry to include only the macros referenced
+    in the provided set of macro references.
+
+    Args:
+        jinja_macros: The original Jinja macro registry containing all macros.
+        package_name: The name of the package for which to create the registry.
+        jinja_references: A set of macro references to retain in the new registry.
+
+    Returns:
+        A new JinjaMacroRegistry containing only the macros for the specified
+        package and the referenced macros.
+    """
+
+    jinja_registry = jinja_macros.copy()
+    jinja_registry.root_macros = jinja_registry.packages.get(package_name) or {}
+    jinja_registry = jinja_registry.trim(jinja_references)
+
+    return jinja_registry