Feat(dbt): Add support for on-run-start and on-run-end hooks

Author: themisvaltinos

docs/integrations/dbt.md

@@ -324,7 +324,6 @@ The dbt jinja methods that are not currently supported are:
 * selected_sources
 * adapter.expand_target_column_types
 * adapter.rename_relation
-* schemas
 * graph.nodes.values
 * graph.metrics.values
 

sqlmesh/core/environment.py

@@ -14,6 +14,7 @@
 from sqlmesh.core.snapshot import SnapshotId, SnapshotTableInfo, Snapshot
 from sqlmesh.utils import word_characters_only
 from sqlmesh.utils.date import TimeLike, now_timestamp
+from sqlmesh.utils.jinja import JinjaMacroRegistry
 from sqlmesh.utils.metaprogramming import Executable
 from sqlmesh.utils.pydantic import PydanticModel, field_validator
 
@@ -218,6 +219,7 @@ class EnvironmentStatements(PydanticModel):
     before_all: t.List[str]
     after_all: t.List[str]
     python_env: t.Dict[str, Executable]
+    jinja_macros: JinjaMacroRegistry = JinjaMacroRegistry()
 
 
 def execute_environment_statements(
@@ -239,6 +241,7 @@ def execute_environment_statements(
             dialect=adapter.dialect,
             default_catalog=default_catalog,
             python_env=statements.python_env,
+            jinja_macros=statements.jinja_macros,
             snapshots=snapshots,
             start=start,
             end=end,

sqlmesh/core/renderer.py

@@ -107,6 +107,15 @@ def _render(
 
         if environment_naming_info := kwargs.get("environment_naming_info", None):
             kwargs["this_env"] = getattr(environment_naming_info, "name")
+            if snapshots and (
+                schemas := set(
+                    [
+                        s.qualified_view_name.schema_for_environment(environment_naming_info)
+                        for s in snapshots.values()
+                    ]
+                )
+            ):
+                kwargs["schemas"] = list(schemas)
 
         this_model = kwargs.pop("this_model", None)
 
@@ -411,19 +420,21 @@ def render(
 
 def render_statements(
     statements: t.List[str],
-    dialect: DialectType = None,
+    dialect: str,
     default_catalog: t.Optional[str] = None,
     python_env: t.Optional[t.Dict[str, Executable]] = None,
+    jinja_macros: t.Optional[JinjaMacroRegistry] = None,
     **render_kwargs: t.Any,
 ) -> t.List[str]:
     rendered_statements: t.List[str] = []
     for statement in statements:
-        for expression in parse(statement, dialect=dialect):
+        for expression in d.parse(statement, default_dialect=dialect):
             if expression:
                 rendered = ExpressionRenderer(
                     expression,
                     dialect,
                     [],
+                    jinja_macro_registry=jinja_macros or JinjaMacroRegistry(),
                     python_env=python_env,
                     default_catalog=default_catalog,
                     quote_identifiers=False,

sqlmesh/dbt/loader.py

@@ -3,6 +3,8 @@
 import logging
 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 (
@@ -11,9 +13,11 @@
     GatewayConfig,
     ModelDefaultsConfig,
 )
+from sqlmesh.core.environment import EnvironmentStatements
 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.context import DbtContext
@@ -23,7 +27,7 @@
 from sqlmesh.dbt.target import TargetConfig
 from sqlmesh.utils import UniqueKeyDict
 from sqlmesh.utils.errors import ConfigError
-from sqlmesh.utils.jinja import JinjaMacroRegistry
+from sqlmesh.utils.jinja import JinjaMacroRegistry, extract_macro_references_and_variables
 
 if sys.version_info >= (3, 12):
     from importlib import metadata
@@ -230,6 +234,58 @@ 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:
+        """Loads dbt's on_run_start, on_run_end hooks into sqlmesh's before_all, after_all statements respectively."""
+
+        on_run_start = []
+        on_run_end = []
+
+        dialect = self.config.dialect
+        for project in self._load_projects():
+            if manifest := project.context._manifest:
+                if stmts := manifest._on_run_start:
+                    on_run_start.extend(stmts)
+                if stmts := manifest._on_run_end:
+                    on_run_end.extend(stmts)
+
+        if statements := on_run_start + on_run_end:
+            jinja_macro_references, used_variables = extract_macro_references_and_variables(
+                *(gen(e) for e in statements)
+            )
+
+            if jinja_macros := project.context.jinja_macros:
+                if root_package := jinja_macros.root_package_name:
+                    jinja_macros.root_macros = jinja_macros.packages[root_package]
+                jinja_macros = (
+                    jinja_macros
+                    if jinja_macros.trimmed
+                    else jinja_macros.trim(jinja_macro_references)
+                )
+            else:
+                jinja_macros = JinjaMacroRegistry()
+
+            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 or macro.get_registry(),
+                variables=self._get_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
+
     def _compute_yaml_max_mtime_per_subfolder(self, root: Path) -> t.Dict[Path, float]:
         if not root.is_dir():
             return {}

sqlmesh/dbt/manifest.py

@@ -94,6 +94,9 @@ 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
+
     def tests(self, package_name: t.Optional[str] = None) -> TestConfigs:
         self._load_all()
         return self._tests_per_package[package_name or self._project_name]
@@ -312,6 +315,11 @@ 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):

tests/core/test_context.py

@@ -1417,6 +1417,7 @@ def test_environment_statements(tmp_path: pathlib.Path):
         after_all=[
             "@grant_schema_usage()",
             "@grant_select_privileges()",
+            "@grant_usage_role(@schemas, 'admin')",
         ],
     )
 
@@ -1481,6 +1482,22 @@ def grant_schema_usage(evaluator):
 """,
     )
 
+    create_temp_file(
+        tmp_path,
+        pathlib.Path(macros_dir, "grant_usage_file.py"),
+        """
+from sqlmesh import macro
+
+@macro()
+def grant_usage_role(evaluator, schemas, role):
+    if evaluator._environment_naming_info:
+        return [
+            f"GRANT USAGE ON SCHEMA {schema} TO {role};"
+            for schema in schemas
+        ]
+""",
+    )
+
     context = Context(paths=tmp_path, config=config)
     snapshots = {s.name: s for s in context.snapshots.values()}
 
@@ -1515,6 +1532,7 @@ def grant_schema_usage(evaluator):
     assert after_all_rendered == [
         "GRANT USAGE ON SCHEMA db TO user_role",
         "GRANT SELECT ON VIEW memory.db.test_after_model TO ROLE admin_role",
+        'GRANT USAGE ON SCHEMA "db" TO "admin"',
     ]
 
     after_all_rendered_dev = render_statements(
@@ -1529,6 +1547,7 @@ def grant_schema_usage(evaluator):
     assert after_all_rendered_dev == [
         "GRANT USAGE ON SCHEMA db__dev TO user_role",
         "GRANT SELECT ON VIEW memory.db__dev.test_after_model TO ROLE admin_role",
+        'GRANT USAGE ON SCHEMA "db__dev" TO "admin"',
     ]
 
 

tests/dbt/test_adapter.py

@@ -13,6 +13,9 @@
 
 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
@@ -270,3 +273,49 @@ def test_quote_as_configured():
     adapter.quote_as_configured("foo", "identifier") == '"foo"'
     adapter.quote_as_configured("foo", "schema") == "foo"
     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) == 1
+    environment_statements = sushi_context._environment_statements[0]
+
+    assert 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 == [
+        "JINJA_STATEMENT_BEGIN;\n{{ create_tables(schemas) }}\nJINJA_END;"
+    ]
+    assert "create_tables" in environment_statements.jinja_macros.root_macros
+
+    rendered_before_all = render_statements(
+        environment_statements.before_all,
+        dialect=sushi_context.default_dialect,
+        python_env=environment_statements.python_env,
+        jinja_macros=environment_statements.jinja_macros,
+        runtime_stage=RuntimeStage.BEFORE_ALL,
+    )
+
+    rendered_after_all = render_statements(
+        environment_statements.after_all,
+        dialect=sushi_context.default_dialect,
+        python_env=environment_statements.python_env,
+        jinja_macros=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)"
+    ]
+
+    # 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_raw__dev AS SELECT 'raw__dev' AS schema",
+            "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",
+        ]
+    )

tests/dbt/test_transformation.py

@@ -997,6 +997,16 @@ def test_dbt_version(sushi_test_project: Project):
     assert context.render("{{ dbt_version }}").startswith("1.")
 
 
+@pytest.mark.xdist_group("dbt_manifest")
+def test_dbt_on_run_start_end(sushi_test_project: Project):
+    context = sushi_test_project.context
+    assert context._manifest
+    assert context._manifest._on_run_start == [
+        "CREATE TABLE IF NOT EXISTS analytic_stats (physical_table VARCHAR, evaluation_time VARCHAR);"
+    ]
+    assert context._manifest._on_run_end == ["{{ create_tables(schemas) }}"]
+
+
 @pytest.mark.xdist_group("dbt_manifest")
 def test_parsetime_adapter_call(
     assert_exp_eq, sushi_test_project: Project, sushi_test_dbt_context: Context

tests/fixtures/dbt/sushi_test/dbt_project.yml

@@ -25,14 +25,14 @@ models:
     +materialized: table
     +pre-hook:
       - '{{ log("pre-hook") }}'
-    +post-hook: 
+    +post-hook:
       - '{{ log("post-hook") }}'
 
 seeds:
   sushi:
     +pre-hook:
       - '{{ log("pre-hook") }}'
-    +post-hook: 
+    +post-hook:
       - '{{ log("post-hook") }}'
 
 vars:
@@ -57,3 +57,9 @@ vars:
       value: 1
     - name: 'item2'
       value: 2
+
+
+on-run-start:
+  - 'CREATE TABLE IF NOT EXISTS analytic_stats (physical_table VARCHAR, evaluation_time VARCHAR);'
+on-run-end:
+  - '{{ create_tables(schemas) }}'

tests/fixtures/dbt/sushi_test/macros/create_tables.sql

@@ -0,0 +1,5 @@
+{% macro create_tables(schemas) %}
+    {% for schema in schemas %}
+        create or replace table schema_table_{{schema}} as select '{{schema}}' as schema;
+    {% endfor%}
+{% endmacro %}