Refactor to handle multiple projects; add multi_repo dbt test

Author: themisvaltinos

examples/multi_dbt/bronze/dbt_project.yml

@@ -19,3 +19,6 @@ require-dbt-version: [">=1.0.0", "<2.0.0"]
 models:
   start: "2024-01-01"
   +materialized: table
+
+on-run-start:
+  - 'CREATE TABLE IF NOT EXISTS analytic_stats (physical_table VARCHAR, evaluation_time VARCHAR);'

examples/multi_dbt/silver/dbt_project.yml

@@ -19,3 +19,6 @@ require-dbt-version: [">=1.0.0", "<2.0.0"]
 models:
   start: "2024-01-01"
   +materialized: table
+
+on-run-end:
+  - '{{ store_schemas(schemas) }}'

examples/multi_dbt/silver/macros/store_schemas.sql

@@ -0,0 +1,3 @@
+{% macro store_schemas(schemas) %}
+    create or replace table schema_table as select {{schemas}} as all_schemas;
+{% endmacro %}

sqlmesh/dbt/loader.py

@@ -27,7 +27,11 @@
 from sqlmesh.dbt.target import TargetConfig
 from sqlmesh.utils import UniqueKeyDict
 from sqlmesh.utils.errors import ConfigError
-from sqlmesh.utils.jinja import JinjaMacroRegistry, extract_macro_references_and_variables
+from sqlmesh.utils.jinja import (
+    JinjaMacroRegistry,
+    MacroInfo,
+    extract_macro_references_and_variables,
+)
 
 if sys.version_info >= (3, 12):
     from importlib import metadata
@@ -237,39 +241,41 @@ def _load_requirements(self) -> t.Tuple[t.Dict[str, str], t.Set[str]]:
     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 = []
-
+        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()
         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)
+            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(e) for e in statements)
+                *(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
             )
-
-            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(),
+                macros=macros,
+                variables=variables,
                 used_variables=used_variables,
                 path=self.config_path,
             )

tests/core/test_integration.py

@@ -4502,6 +4502,24 @@ def test_multi_dbt(mocker):
     context.apply(plan)
     validate_apply_basics(context, c.PROD, plan.snapshots.values())
 
+    environment_statements = context.state_sync.get_environment_statements(c.PROD)
+    assert len(environment_statements) == 2
+    bronze_statements = environment_statements[0]
+    assert bronze_statements.before_all == [
+        "JINJA_STATEMENT_BEGIN;\nCREATE TABLE IF NOT EXISTS analytic_stats (physical_table VARCHAR, evaluation_time VARCHAR);\nJINJA_END;"
+    ]
+    assert not bronze_statements.after_all
+    silver_statements = environment_statements[1]
+    assert not silver_statements.before_all
+    assert silver_statements.after_all == [
+        "JINJA_STATEMENT_BEGIN;\n{{ store_schemas(schemas) }}\nJINJA_END;"
+    ]
+    assert "store_schemas" in silver_statements.jinja_macros.root_macros
+    analytics_table = context.fetchdf("select * from analytic_stats;")
+    assert sorted(analytics_table.columns) == sorted(["physical_table", "evaluation_time"])
+    schema_table = context.fetchdf("select * from schema_table;")
+    assert sorted(schema_table.all_schemas[0]) == sorted(["bronze", "silver"])
+
 
 def test_multi_hybrid(mocker):
     context = Context(