183 lines
7.2 KiB
Python
183 lines
7.2 KiB
Python
from typing import (
|
|
Any, Dict, Iterable, Union, Optional, List, Iterator, Mapping, Set
|
|
)
|
|
|
|
from dbt.clients.jinja import MacroGenerator, MacroStack
|
|
from dbt.contracts.graph.parsed import ParsedMacro
|
|
from dbt.include.global_project import PROJECT_NAME as GLOBAL_PROJECT_NAME
|
|
from dbt.exceptions import (
|
|
raise_duplicate_macro_name, raise_compiler_error
|
|
)
|
|
|
|
|
|
FlatNamespace = Dict[str, MacroGenerator]
|
|
NamespaceMember = Union[FlatNamespace, MacroGenerator]
|
|
FullNamespace = Dict[str, NamespaceMember]
|
|
|
|
|
|
# The point of this class is to collect the various macros
|
|
# and provide the ability to flatten them into the ManifestContexts
|
|
# that are created for jinja, so that macro calls can be resolved.
|
|
# Creates special iterators and _keys methods to flatten the lists.
|
|
# When this class is created it has a static 'local_namespace' which
|
|
# depends on the package of the node, so it only works for one
|
|
# particular local package at a time for "flattening" into a context.
|
|
# 'get_by_package' should work for any macro.
|
|
class MacroNamespace(Mapping):
|
|
def __init__(
|
|
self,
|
|
global_namespace: FlatNamespace, # root package macros
|
|
local_namespace: FlatNamespace, # packages for *this* node
|
|
global_project_namespace: FlatNamespace, # internal packages
|
|
packages: Dict[str, FlatNamespace], # non-internal packages
|
|
):
|
|
self.global_namespace: FlatNamespace = global_namespace
|
|
self.local_namespace: FlatNamespace = local_namespace
|
|
self.packages: Dict[str, FlatNamespace] = packages
|
|
self.global_project_namespace: FlatNamespace = global_project_namespace
|
|
|
|
def _search_order(self) -> Iterable[Union[FullNamespace, FlatNamespace]]:
|
|
yield self.local_namespace # local package
|
|
yield self.global_namespace # root package
|
|
yield self.packages # non-internal packages
|
|
yield {
|
|
GLOBAL_PROJECT_NAME: self.global_project_namespace, # dbt
|
|
}
|
|
yield self.global_project_namespace # other internal project besides dbt
|
|
|
|
# provides special keys method for MacroNamespace iterator
|
|
# returns keys from local_namespace, global_namespace, packages,
|
|
# global_project_namespace
|
|
def _keys(self) -> Set[str]:
|
|
keys: Set[str] = set()
|
|
for search in self._search_order():
|
|
keys.update(search)
|
|
return keys
|
|
|
|
# special iterator using special keys
|
|
def __iter__(self) -> Iterator[str]:
|
|
for key in self._keys():
|
|
yield key
|
|
|
|
def __len__(self):
|
|
return len(self._keys())
|
|
|
|
def __getitem__(self, key: str) -> NamespaceMember:
|
|
for dct in self._search_order():
|
|
if key in dct:
|
|
return dct[key]
|
|
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:
|
|
return self.global_project_namespace.get(name)
|
|
elif package_name in self.packages:
|
|
return self.packages[package_name].get(name)
|
|
else:
|
|
raise_compiler_error(
|
|
f"Could not find package '{package_name}'"
|
|
)
|
|
|
|
|
|
# This class builds the MacroNamespace by adding macros to
|
|
# internal_packages or packages, and locals/globals.
|
|
# Call 'build_namespace' to return a MacroNamespace.
|
|
# This is used by ManifestContext (and subclasses)
|
|
class MacroNamespaceBuilder:
|
|
def __init__(
|
|
self,
|
|
root_package: str,
|
|
search_package: str,
|
|
thread_ctx: MacroStack,
|
|
internal_packages: List[str],
|
|
node: Optional[Any] = None,
|
|
) -> None:
|
|
self.root_package = root_package
|
|
self.search_package = search_package
|
|
# internal packages comes from get_adapter_package_names
|
|
self.internal_package_names = set(internal_packages)
|
|
self.internal_package_names_order = internal_packages
|
|
# macro_func is added here if in root package, since
|
|
# the root package acts as a "global" namespace, overriding
|
|
# everything else except local external package macro calls
|
|
self.globals: FlatNamespace = {}
|
|
# macro_func is added here if it's the package for this node
|
|
self.locals: FlatNamespace = {}
|
|
# Create a dictionary of [package name][macro name] =
|
|
# MacroGenerator object which acts like a function
|
|
self.internal_packages: Dict[str, FlatNamespace] = {}
|
|
self.packages: Dict[str, FlatNamespace] = {}
|
|
self.thread_ctx = thread_ctx
|
|
self.node = node
|
|
|
|
def _add_macro_to(
|
|
self,
|
|
hierarchy: Dict[str, FlatNamespace],
|
|
macro: ParsedMacro,
|
|
macro_func: MacroGenerator,
|
|
):
|
|
if macro.package_name in hierarchy:
|
|
namespace = hierarchy[macro.package_name]
|
|
else:
|
|
namespace = {}
|
|
hierarchy[macro.package_name] = namespace
|
|
|
|
if macro.name in namespace:
|
|
raise_duplicate_macro_name(
|
|
macro_func.macro, macro, macro.package_name
|
|
)
|
|
hierarchy[macro.package_name][macro.name] = macro_func
|
|
|
|
def add_macro(self, macro: ParsedMacro, ctx: Dict[str, Any]):
|
|
macro_name: str = macro.name
|
|
|
|
# MacroGenerator is in clients/jinja.py
|
|
# a MacroGenerator object is a callable object that will
|
|
# execute the MacroGenerator.__call__ function
|
|
macro_func: MacroGenerator = MacroGenerator(
|
|
macro, ctx, self.node, self.thread_ctx
|
|
)
|
|
|
|
# internal macros (from plugins) will be processed separately from
|
|
# project macros, so store them in a different place
|
|
if macro.package_name in self.internal_package_names:
|
|
self._add_macro_to(self.internal_packages, macro, macro_func)
|
|
else:
|
|
# if it's not an internal package
|
|
self._add_macro_to(self.packages, macro, macro_func)
|
|
# add to locals if it's the package this node is in
|
|
if macro.package_name == self.search_package:
|
|
self.locals[macro_name] = macro_func
|
|
# add to globals if it's in the root package
|
|
elif macro.package_name == self.root_package:
|
|
self.globals[macro_name] = macro_func
|
|
|
|
def add_macros(self, macros: Iterable[ParsedMacro], ctx: Dict[str, Any]):
|
|
for macro in macros:
|
|
self.add_macro(macro, ctx)
|
|
|
|
def build_namespace(
|
|
self, macros: Iterable[ParsedMacro], ctx: Dict[str, Any]
|
|
) -> MacroNamespace:
|
|
self.add_macros(macros, ctx)
|
|
|
|
# Iterate in reverse-order and overwrite: the packages that are first
|
|
# in the list are the ones we want to "win".
|
|
global_project_namespace: FlatNamespace = {}
|
|
for pkg in reversed(self.internal_package_names_order):
|
|
if pkg in self.internal_packages:
|
|
# add the macros pointed to by this package name
|
|
global_project_namespace.update(self.internal_packages[pkg])
|
|
|
|
return MacroNamespace(
|
|
global_namespace=self.globals, # root package macros
|
|
local_namespace=self.locals, # packages for *this* node
|
|
global_project_namespace=global_project_namespace, # internal packages
|
|
packages=self.packages, # non internal_packages
|
|
)
|