add macros to flat graph

dbt/clients/jinja.py

@@ -4,36 +4,56 @@
 import jinja2
 import jinja2.sandbox
 
+from dbt.utils import NodeType
 
-class SilentUndefined(jinja2.Undefined):
-    """
-    This class sets up the parser to just ignore undefined jinja2 calls. So,
-    for example, `env` is not defined here, but will not make the parser fail
-    with a fatal error.
-    """
-    def _fail_with_undefined_error(self, *args, **kwargs):
-        return None
 
-    __add__ = __radd__ = __mul__ = __rmul__ = __div__ = __rdiv__ = \
-        __truediv__ = __rtruediv__ = __floordiv__ = __rfloordiv__ = \
-        __mod__ = __rmod__ = __pos__ = __neg__ = __call__ = \
-        __getitem__ = __lt__ = __le__ = __gt__ = __ge__ = __int__ = \
-        __float__ = __complex__ = __pow__ = __rpow__ = \
-        _fail_with_undefined_error
+def create_macro_capture_env(node):
 
+    class ParserMacroCapture(jinja2.Undefined):
+        """
+        This class sets up the parser to capture macros.
+        """
+        def __init__(self, hint=None, obj=None, name=None,
+                     exc=None):
+            super(jinja2.Undefined, self).__init__()
 
-env = jinja2.sandbox.SandboxedEnvironment()
+            self.node = node
+            self.name = name
+            self.package_name = node.get('package_name')
+
+        def __getattr__(self, name):
+
+            # jinja uses these for safety, so we have to override them.
+            # see https://github.com/pallets/jinja/blob/master/jinja2/sandbox.py#L332-L339 # noqa
+            if name in ['unsafe_callable', 'alters_data']:
+                return False
+
+            self.package_name = self.name
+            self.name = name
+
+            return self
+
+        def __call__(self, *args, **kwargs):
+            path = '{}.{}.{}'.format(NodeType.Macro,
+                                     self.package_name,
+                                     self.name)
 
-silent_on_undefined_env = jinja2.sandbox.SandboxedEnvironment(
-    undefined=SilentUndefined)
+            if path not in self.node['depends_on']['macros']:
+                self.node['depends_on']['macros'].append(path)
 
+    return jinja2.sandbox.SandboxedEnvironment(
+        undefined=ParserMacroCapture)
 
-def get_template(string, ctx, node=None, silent_on_undefined=False):
+
+env = jinja2.sandbox.SandboxedEnvironment()
+
+
+def get_template(string, ctx, node=None, capture_macros=False):
     try:
         local_env = env
 
-        if silent_on_undefined:
-            local_env = silent_on_undefined_env
+        if capture_macros is True:
+            local_env = create_macro_capture_env(node)
 
         return local_env.from_string(dbt.compat.to_string(string), globals=ctx)
 
@@ -42,11 +62,15 @@ def get_template(string, ctx, node=None, silent_on_undefined=False):
         dbt.exceptions.raise_compiler_error(node, str(e))
 
 
-def get_rendered(string, ctx, node=None, silent_on_undefined=False):
+def render_template(template, ctx, node=None):
     try:
-        template = get_template(string, ctx, node, silent_on_undefined)
         return template.render(ctx)
 
     except (jinja2.exceptions.TemplateSyntaxError,
             jinja2.exceptions.UndefinedError) as e:
         dbt.exceptions.raise_compiler_error(node, str(e))
+
+
+def get_rendered(string, ctx, node=None, capture_macros=False):
+    template = get_template(string, ctx, node, capture_macros)
+    return render_template(template, ctx, node=None)