feat: support custom macros in dbt vars (#5250)

Author: eakmanrq

sqlmesh/dbt/context.py

@@ -101,7 +101,10 @@ def add_variables(self, variables: t.Dict[str, t.Any]) -> None:
         self._jinja_environment = None
 
     def set_and_render_variables(self, variables: t.Dict[str, t.Any], package: str) -> None:
-        jinja_environment = self.jinja_macros.build_environment(**self.jinja_globals)
+        package_macros = self.jinja_macros.copy(
+            update={"top_level_packages": [*self.jinja_macros.top_level_packages, package]}
+        )
+        jinja_environment = package_macros.build_environment(**self.jinja_globals)
 
         def _render_var(value: t.Any) -> t.Any:
             if isinstance(value, str):

tests/dbt/test_config.py

@@ -8,6 +8,8 @@
 from pytest_mock import MockerFixture
 
 from sqlglot import exp
+
+from sqlmesh import Context
 from sqlmesh.core.audit import StandaloneAudit
 from sqlmesh.core.config import Config, ModelDefaultsConfig
 from sqlmesh.core.dialect import jinja_query
@@ -360,7 +362,7 @@ def test_variables(assert_exp_eq, sushi_test_project):
     # Finally, check that variable scoping & overwriting (some_var) works as expected
     expected_sushi_variables = {
         "yet_another_var": 1,
-        "top_waiters:limit": 10,
+        "top_waiters:limit": "{{ get_top_waiters_limit() }}",
         "top_waiters:revenue": "revenue",
         "customers:boo": ["a", "b"],
         "nested_vars": {
@@ -389,6 +391,11 @@ def test_variables(assert_exp_eq, sushi_test_project):
     assert sushi_test_project.packages["customers"].variables == expected_customer_variables
 
 
+@pytest.mark.slow
+def test_jinja_in_dbt_variables(sushi_test_dbt_context: Context):
+    assert sushi_test_dbt_context.render("sushi.top_waiters").sql().endswith("LIMIT 10")
+
+
 @pytest.mark.slow
 def test_nested_variables(sushi_test_project):
     model_config = ModelConfig(

tests/dbt/test_manifest.py

@@ -34,7 +34,11 @@ def test_manifest_helper(caplog):
         refs={"sushi.waiter_revenue_by_day", "waiter_revenue_by_day"},
         variables={"top_waiters:revenue", "top_waiters:limit"},
         model_attrs={"columns", "config"},
-        macros=[MacroReference(name="ref"), MacroReference(name="var")],
+        macros=[
+            MacroReference(name="get_top_waiters_limit"),
+            MacroReference(name="ref"),
+            MacroReference(name="var"),
+        ],
     )
     assert models["top_waiters"].materialized == "view"
     assert models["top_waiters"].dialect_ == "postgres"
@@ -181,7 +185,7 @@ def test_variable_override():
         profile.target,
         model_defaults=ModelDefaultsConfig(start="2020-01-01"),
     )
-    assert helper.models()["top_waiters"].limit_value == 10
+    assert helper.models()["top_waiters"].limit_value.strip() == "10"
 
     helper = ManifestHelper(
         project_path,

tests/fixtures/dbt/sushi_test/dbt_project.yml

@@ -39,7 +39,7 @@ sources:
     identifier: false
 
 vars:
-  top_waiters:limit: 10
+  top_waiters:limit: "{{ get_top_waiters_limit() }}"
   'top_waiters:revenue': "revenue"
 
   # The following are only used for testing purposes

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

@@ -0,0 +1,3 @@
+{% macro get_top_waiters_limit() %}
+10
+{% endmacro %}