Refactor macro contexts.

.changes/unreleased/Under the Hood-20240122-165446.yaml

@@ -0,0 +1,6 @@
+kind: Under the Hood
+body: Clean up macro contexts.
+time: 2024-01-22T16:54:46.247418-05:00
+custom:
+  Author: peterallenwebb
+  Issue: "9422"

core/dbt/context/macros.py

@@ -67,7 +67,6 @@ def __getitem__(self, key: str) -> NamespaceMember:
         raise KeyError(key)
 
     def get_from_package(self, package_name: Optional[str], name: str) -> Optional[MacroGenerator]:
-        pkg: FlatNamespace
         if package_name is None:
             return self.get(name)
         elif package_name == GLOBAL_PROJECT_NAME:
@@ -125,7 +124,7 @@ def _add_macro_to(
             raise DuplicateMacroNameError(macro_func.macro, macro, macro.package_name)
         hierarchy[macro.package_name][macro.name] = macro_func
 
-    def add_macro(self, macro: Macro, ctx: Dict[str, Any]):
+    def add_macro(self, macro: Macro, ctx: Dict[str, Any]) -> None:
         macro_name: str = macro.name
 
         # MacroGenerator is in clients/jinja.py
@@ -147,12 +146,15 @@ def add_macro(self, macro: Macro, ctx: Dict[str, Any]):
             elif macro.package_name == self.root_package:
                 self.globals[macro_name] = macro_func
 
-    def add_macros(self, macros: Iterable[Macro], ctx: Dict[str, Any]):
+    def add_macros(self, macros: Iterable[Macro], ctx: Dict[str, Any]) -> None:
         for macro in macros:
             self.add_macro(macro, ctx)
 
-    def build_namespace(self, macros: Iterable[Macro], ctx: Dict[str, Any]) -> MacroNamespace:
-        self.add_macros(macros, ctx)
+    def build_namespace(
+        self, macros_by_package: Dict[str, Dict[str, Macro]], ctx: Dict[str, Any]
+    ) -> MacroNamespace:
+        for package in macros_by_package.values():
+            self.add_macros(package.values(), ctx)
 
         # Iterate in reverse-order and overwrite: the packages that are first
         # in the list are the ones we want to "win".

core/dbt/context/manifest.py

@@ -8,7 +8,7 @@
 
 
 from .configured import ConfiguredContext
-from .macros import MacroNamespaceBuilder
+from .macros import MacroNamespace, MacroNamespaceBuilder
 
 
 class ManifestContext(ConfiguredContext):
@@ -36,11 +36,11 @@ def __init__(
         # to be able to do: namespace.get_from_package(..)
         self.namespace = self._build_namespace()
 
-    def _build_namespace(self):
+    def _build_namespace(self) -> MacroNamespace:
         # this takes all the macros in the manifest and adds them
         # to the MacroNamespaceBuilder stored in self.namespace
         builder = self._get_namespace_builder()
-        return builder.build_namespace(self.manifest.macros.values(), self._ctx)
+        return builder.build_namespace(self.manifest.get_macros_by_package(), self._ctx)
 
     def _get_namespace_builder(self) -> MacroNamespaceBuilder:
         # avoid an import loop
@@ -65,6 +65,7 @@ def to_dict(self):
             dct.update(self.namespace.project_namespace)
         else:
             dct.update(self.namespace)
+
         return dct
 
     @contextproperty()

core/dbt/contracts/graph/manifest.py

@@ -689,6 +689,8 @@
     def __init__(self):
         self.macros = []
         self.metadata = {}
+        self._macros_by_name = {}
+        self._macros_by_package = {}
 
     def find_macro_by_name(
         self, name: str, root_project_name: str, package: Optional[str]
@@ -757,10 +759,13 @@
         from dbt.adapters.factory import get_adapter_package_names
 
         candidates: CandidateList = CandidateList()
+
+        macros_by_name = self.get_macros_by_name()
+        if name not in macros_by_name:
+            return candidates
+
         packages = set(get_adapter_package_names(self.metadata.adapter_type))
-        for unique_id, macro in self.macros.items():
-            if macro.name != name:
-                continue
+        for macro in macros_by_name[name]:
             candidate = MacroCandidate(
                 locality=_get_locality(macro, root_project_name, packages),
                 macro=macro,
@@ -770,6 +775,49 @@
 
         return candidates
 
+    def get_macros_by_name(self) -> Dict[str, List[Macro]]:
+        if self._macros_by_name is None:
+            # The by-name mapping doesn't exist yet (perhaps because the manifest
+            # was deserialized), so we build it.
+            self._macros_by_name = self._build_macros_by_name(self.macros)
+
+        return self._macros_by_name
+
+    @staticmethod
+    def _build_macros_by_name(macros: Mapping[str, Macro]) -> Dict[str, List[Macro]]:
+        # Convert a macro dictionary keyed on unique id to a flattened version
+        # keyed on macro name for faster lookup by name. Since macro names are
+        # not necessarily unique, the dict value is a list.
+        macros_by_name: Dict[str, List[Macro]] = {}
+        for macro in macros.values():
+            if macro.name not in macros_by_name:
+                macros_by_name[macro.name] = []
+
+            macros_by_name[macro.name].append(macro)
+
+        return macros_by_name
+
+    def get_macros_by_package(self) -> Dict[str, Dict[str, Macro]]:
+        if self._macros_by_package is None:
+            # The by-package mapping doesn't exist yet (perhaps because the manifest
+            # was deserialized), so we build it.
+            self._macros_by_package = self._build_macros_by_package(self.macros)
+
+        return self._macros_by_package
+
+    @staticmethod
+    def _build_macros_by_package(macros: Mapping[str, Macro]) -> Dict[str, Dict[str, Macro]]:
+        # Convert a macro dictionary keyed on unique id to a flattened version
+        # keyed on package name for faster lookup by name.
+        macros_by_package: Dict[str, Dict[str, Macro]] = {}
+        for macro in macros.values():
+            if macro.package_name not in macros_by_package:
+                macros_by_package[macro.package_name] = {}
+            macros_by_name = macros_by_package[macro.package_name]
+            macros_by_name[macro.name] = macro
+
+        return macros_by_package
+
 
 @dataclass
 class ParsingInfo:
@@ -845,6 +893,14 @@
         default_factory=get_mp_context().Lock,
         metadata={"serialize": lambda x: None, "deserialize": lambda x: None},
     )
+    _macros_by_name: Optional[Dict[str, List[Macro]]] = field(
+        default=None,
+        metadata={"serialize": lambda x: None, "deserialize": lambda x: None},
+    )
+    _macros_by_package: Optional[Dict[str, Dict[str, Macro]]] = field(
+        default=None,
+        metadata={"serialize": lambda x: None, "deserialize": lambda x: None},
+    )
 
     def __pre_serialize__(self):
         # serialization won't work with anything except an empty source_patches because
@@ -1408,6 +1464,23 @@
             raise DuplicateMacroInPackageError(macro=macro, macro_mapping=self.macros)
 
         self.macros[macro.unique_id] = macro
+
+        if self._macros_by_name is None:
+            self._macros_by_name = self._build_macros_by_name(self.macros)
+
+        if macro.name not in self._macros_by_name:
+            self._macros_by_name[macro.name] = []
+
+        self._macros_by_name[macro.name].append(macro)
+
+        if self._macros_by_package is None:
+            self._macros_by_package = self._build_macros_by_package(self.macros)
+
+        if macro.package_name not in self._macros_by_package:
+            self._macros_by_package[macro.package_name] = {}
+
+        self._macros_by_package[macro.package_name][macro.name] = macro
+
         source_file.macros.append(macro.unique_id)
 
     def has_file(self, source_file: SourceFile) -> bool:
@@ -1565,6 +1638,8 @@
         # This is returned by the 'graph' context property
         # in the ProviderContext class.
         self.flat_graph: Dict[str, Any] = {}
+        self._macros_by_name: Optional[Dict[str, List[Macro]]] = None
+        self._macros_by_package: Optional[Dict[str, Dict[str, Macro]]] = None
 
 
 AnyManifest = Union[Manifest, MacroManifest]

tests/unit/test_context.py

@@ -1,4 +1,3 @@
-import itertools
 import unittest
 import os
 from typing import Set, Dict, Any
@@ -317,10 +316,21 @@ def mock_macro(name, package_name):
 
 def mock_manifest(config):
     manifest_macros = {}
+    macros_by_package = {}
     for name in ["macro_a", "macro_b"]:
         macro = mock_macro(name, config.project_name)
         manifest_macros[macro.unique_id] = macro
-    return mock.MagicMock(macros=manifest_macros)
+        if macro.package_name not in macros_by_package:
+            macros_by_package[macro.package_name] = {}
+        macro_package = macros_by_package[macro.package_name]
+        macro_package[macro.name] = macro
+
+    def gmbp():
+        return macros_by_package
+
+    m = mock.MagicMock(macros=manifest_macros)
+    m.get_macros_by_package = gmbp
+    return m
 
 
 def mock_model():
@@ -469,15 +479,19 @@ def test_macro_namespace_duplicates(config_postgres, manifest_fx):
 def test_macro_namespace(config_postgres, manifest_fx):
     mn = macros.MacroNamespaceBuilder("root", "search", MacroStack(), ["dbt_postgres", "dbt"])
 
+    mbp = manifest_fx.get_macros_by_package()
     dbt_macro = mock_macro("some_macro", "dbt")
+    mbp["dbt"] = {"some_macro": dbt_macro}
+
     # same namespace, same name, different pkg!
     pg_macro = mock_macro("some_macro", "dbt_postgres")
+    mbp["dbt_postgres"] = {"some_macro": pg_macro}
+
     # same name, different package
     package_macro = mock_macro("some_macro", "root")
+    mbp["root"]["some_macro"] = package_macro
 
-    all_macros = itertools.chain(manifest_fx.macros.values(), [dbt_macro, pg_macro, package_macro])
-
-    namespace = mn.build_namespace(all_macros, {})
+    namespace = mn.build_namespace(mbp, {})
     dct = dict(namespace)
     for result in [dct, namespace]:
         assert "dbt" in result