Fix: Take the current package into account when trimming jinja macros (#1010)

Author: izeigerman

sqlmesh/dbt/basemodel.py

@@ -210,7 +210,9 @@ def model_function(self) -> AttributeDict[str, t.Any]:
     def sqlmesh_model_kwargs(self, context: DbtContext) -> t.Dict[str, t.Any]:
         """Get common sqlmesh model parameters"""
         model_context = context.context_for_dependencies(self.dependencies)
-        jinja_macros = model_context.jinja_macros.trim(self.dependencies.macros)
+        jinja_macros = model_context.jinja_macros.trim(
+            self.dependencies.macros, package=self.package_name
+        )
         jinja_macros.global_objs.update(
             {
                 "this": self.relation_info,

sqlmesh/dbt/loader.py

@@ -132,10 +132,7 @@ def _load_project(self) -> Project:
             )
 
         for package_name, macro_infos in context.manifest.all_macros.items():
-            context.jinja_macros.add_macros(
-                macro_infos,
-                package=package_name if package_name != context.project_name else None,
-            )
+            context.jinja_macros.add_macros(macro_infos, package=package_name)
 
         self._macros_max_mtime = max(macros_mtimes) if macros_mtimes else None
 

sqlmesh/dbt/test.py

@@ -35,6 +35,7 @@ class TestConfig(GeneralConfig):
         column_name: The name of the column under test.
         dependencies: The macros, refs, and sources the test depends upon.
         dialect: The sql dialect of the test.
+        package_name: Name of the package that defines the test.
         alias: The alias for the materialized table where failures are stored (Not supported).
         schema: The schema for the materialized table where the failures are stored (Not supported).
         database: The database for the materilized table where the failures are stored (Not supported).
@@ -58,6 +59,7 @@ class TestConfig(GeneralConfig):
     dialect: str = ""
 
     # dbt fields
+    package_name: str = ""
     alias: t.Optional[str] = None
     schema_: t.Optional[str] = Field("", alias="schema")
     database: t.Optional[str] = None
@@ -89,7 +91,9 @@ def to_sqlmesh(self, context: DbtContext) -> Audit:
         """
         test_context = context.context_for_dependencies(self.dependencies)
 
-        jinja_macros = test_context.jinja_macros.trim(self.dependencies.macros)
+        jinja_macros = test_context.jinja_macros.trim(
+            self.dependencies.macros, package=self.package_name
+        )
         jinja_macros.global_objs.update(
             {
                 "config": self.attribute_dict,

sqlmesh/utils/jinja.py

@@ -251,24 +251,31 @@ def build_environment(self, **kwargs: t.Any) -> Environment:
         env.filters.update(self._environment.filters)
         return env
 
-    def trim(self, dependencies: t.Iterable[MacroReference]) -> JinjaMacroRegistry:
+    def trim(
+        self, dependencies: t.Iterable[MacroReference], package: t.Optional[str] = None
+    ) -> JinjaMacroRegistry:
         """Trims the registry by keeping only macros with given references and their transitive dependencies.
 
         Args:
             dependencies: References to macros that should be kept.
+            package: The name of the package in the context of which the trimming should be performed.
 
         Returns:
             A new trimmed registry.
         """
         dependencies_by_package: t.Dict[t.Optional[str], t.Set[str]] = defaultdict(set)
         for dep in dependencies:
-            dependencies_by_package[dep.package].add(dep.name)
+            dependencies_by_package[dep.package or package].add(dep.name)
+
+        top_level_packages = self.top_level_packages.copy()
+        if package is not None:
+            top_level_packages.append(package)
 
         result = JinjaMacroRegistry(
             global_objs=self.global_objs.copy(),
             create_builtins_module=self.create_builtins_module,
             root_package_name=self.root_package_name,
-            top_level_packages=self.top_level_packages.copy(),
+            top_level_packages=top_level_packages,
         )
         for package, names in dependencies_by_package.items():
             result = result.merge(self._trim_macros(names, package))

tests/dbt/conftest.py

@@ -16,10 +16,7 @@ def sushi_test_project() -> Project:
     delete_cache(project_root)
     project = Project.load(DbtContext(project_root=Path(project_root)))
     for package_name, package in project.packages.items():
-        project.context.jinja_macros.add_macros(
-            package.macro_infos,
-            package=package_name if package_name != project.context.project_name else None,
-        )
+        project.context.jinja_macros.add_macros(package.macro_infos, package=package_name)
     return project
 
 

tests/utils/test_jinja.py

@@ -159,6 +159,14 @@ def test_macro_registry_trim():
     )
     assert rendered == "macro_a_a macro_a_b"
 
+    trimmed_registry_for_package_b = registry.trim(
+        [MacroReference(name="macro_b_b")], package="package_b"
+    )
+    assert set(trimmed_registry_for_package_b.packages) == {"package_a", "package_b"}
+    assert set(trimmed_registry_for_package_b.packages["package_a"]) == {"macro_a_a"}
+    assert set(trimmed_registry_for_package_b.packages["package_b"]) == {"macro_b_a", "macro_b_b"}
+    assert not trimmed_registry_for_package_b.root_macros
+
 
 def test_macro_return():
     macros = "{% macro test_return() %}{{ macro_return([1, 2, 3]) }}{% endmacro %}"