In schema tests, pass rendered items (#2149)

CHANGELOG.md

@@ -1,16 +1,16 @@
 ## dbt next (release TBD)
 ### Features
-- Added --fail-fast argument for dbt run and dbt test to fail on first test failure or runtime error. ([#1649](https://github.com/fishtown-analytics/dbt/issues/1649))
-- Support for appending query comments to SQL queries. ([#2138](https://github.com/fishtown-analytics/dbt/issues/2138))
+- Added --fail-fast argument for dbt run and dbt test to fail on first test failure or runtime error. ([#1649](https://github.com/fishtown-analytics/dbt/issues/1649), [#2224](https://github.com/fishtown-analytics/dbt/pull/2224))
+- Support for appending query comments to SQL queries. ([#2138](https://github.com/fishtown-analytics/dbt/issues/2138), [#2199](https://github.com/fishtown-analytics/dbt/pull/2199))
 - Added a `get-manifest` API call. ([#2168](https://github.com/fishtown-analytics/dbt/issues/2168), [#2232](https://github.com/fishtown-analytics/dbt/pull/2232))
+- Users can now use jinja as arguments to tests. Test arguments are rendered in the native context and injected into the test execution context directly. ([#2149](https://github.com/fishtown-analytics/dbt/issues/2149), [#2220](https://github.com/fishtown-analytics/dbt/pull/2220))
 
 ### Fixes
 - When a jinja value is undefined, give a helpful error instead of failing with cryptic "cannot pickle ParserMacroCapture" errors ([#2110](https://github.com/fishtown-analytics/dbt/issues/2110), [#2184](https://github.com/fishtown-analytics/dbt/pull/2184))
 - Added timeout to registry download call ([#2195](https://github.com/fishtown-analytics/dbt/issues/2195), [#2228](https://github.com/fishtown-analytics/dbt/pull/2228))
 
 Contributors:
  - [@raalsky](https://github.com/Raalsky) ([#2224](https://github.com/fishtown-analytics/dbt/pull/2224), [#2228](https://github.com/fishtown-analytics/dbt/pull/2228))
- - [@raalsky](https://github.com/Raalsky) ([#2224](https://github.com/fishtown-analytics/dbt/pull/2224))
  - [@ilkinulas](https://github.com/ilkinulas) [#2199](https://github.com/fishtown-analytics/dbt/pull/2199)
 
 

core/dbt/clients/jinja.py

@@ -1,6 +1,7 @@
 import codecs
 import linecache
 import os
+import re
 import tempfile
 import threading
 from contextlib import contextmanager
@@ -10,15 +11,19 @@
 
 import jinja2
 import jinja2.ext
+import jinja2.nativetypes  # type: ignore
 import jinja2.nodes
 import jinja2.parser
 import jinja2.sandbox
 
 from dbt.utils import (
-    get_dbt_macro_name, get_docs_macro_name, get_materialization_macro_name
+    get_dbt_macro_name, get_docs_macro_name, get_materialization_macro_name,
+    deep_map
 )
 
 from dbt.clients._jinja_blocks import BlockIterator, BlockData, BlockTag
+from dbt.contracts.graph.compiled import CompiledSchemaTestNode
+from dbt.contracts.graph.parsed import ParsedSchemaTestNode
 from dbt.exceptions import (
     InternalException, raise_compiler_error, CompilationException,
     invalid_materialization_argument, MacroReturn
@@ -93,6 +98,33 @@ def _compile(self, source, filename):
         return super()._compile(source, filename)  # type: ignore
 
 
+class NativeSandboxEnvironment(MacroFuzzEnvironment):
+    code_generator_class = jinja2.nativetypes.NativeCodeGenerator
+
+
+class NativeSandboxTemplate(jinja2.nativetypes.NativeTemplate):  # mypy: ignore
+    environment_class = NativeSandboxEnvironment
+
+    def render(self, *args, **kwargs):
+        """Render the template to produce a native Python type. If the
+        result is a single node, its value is returned. Otherwise, the
+        nodes are concatenated as strings. If the result can be parsed
+        with :func:`ast.literal_eval`, the parsed value is returned.
+        Otherwise, the string is returned.
+        """
+        vars = dict(*args, **kwargs)
+        try:
+            return jinja2.nativetypes.native_concat(
+                self.root_render_func(self.new_context(vars)),
+                preserve_quotes=True
+            )
+        except Exception:
+            return self.environment.handle_exception()
+
+
+NativeSandboxEnvironment.template_class = NativeSandboxTemplate  # type: ignore
+
+
 class TemplateCache:
 
     def __init__(self):
@@ -348,7 +380,11 @@ def __reduce__(self):
     return Undefined
 
 
-def get_environment(node=None, capture_macros=False):
+def get_environment(
+    node=None,
+    capture_macros: bool = False,
+    native: bool = False,
+) -> jinja2.Environment:
     args: Dict[str, List[Union[str, Type[jinja2.ext.Extension]]]] = {
         'extensions': ['jinja2.ext.do']
     }
@@ -359,7 +395,13 @@ def get_environment(node=None, capture_macros=False):
     args['extensions'].append(MaterializationExtension)
     args['extensions'].append(DocumentationExtension)
 
-    return MacroFuzzEnvironment(**args)
+    env_cls: Type[jinja2.Environment]
+    if native:
+        env_cls = NativeSandboxEnvironment
+    else:
+        env_cls = MacroFuzzEnvironment
+
+    return env_cls(**args)
 
 
 @contextmanager
@@ -378,21 +420,39 @@ def parse(string):
         return get_environment().parse(str(string))
 
 
-def get_template(string, ctx, node=None, capture_macros=False):
+def get_template(
+    string: str,
+    ctx: Dict[str, Any],
+    node=None,
+    capture_macros: bool = False,
+    native: bool = False,
+):
     with catch_jinja(node):
-        env = get_environment(node, capture_macros)
+        env = get_environment(node, capture_macros, native=native)
 
         template_source = str(string)
         return env.from_string(template_source, globals=ctx)
 
 
-def render_template(template, ctx, node=None):
+def render_template(template, ctx: Dict[str, Any], node=None) -> str:
     with catch_jinja(node):
         return template.render(ctx)
 
 
-def get_rendered(string, ctx, node=None, capture_macros=False):
-    template = get_template(string, ctx, node, capture_macros=capture_macros)
+def get_rendered(
+    string: str,
+    ctx: Dict[str, Any],
+    node=None,
+    capture_macros: bool = False,
+    native: bool = False,
+) -> str:
+    template = get_template(
+        string,
+        ctx,
+        node,
+        capture_macros=capture_macros,
+        native=native,
+    )
 
     return render_template(template, ctx, node)
 
@@ -424,3 +484,35 @@ def extract_toplevel_blocks(
         allowed_blocks=allowed_blocks,
         collect_raw_data=collect_raw_data
     )
+
+
+SCHEMA_TEST_KWARGS_NAME = '_dbt_schema_test_kwargs'
+
+
+def add_rendered_test_kwargs(
+    context: Dict[str, Any],
+    node: Union[ParsedSchemaTestNode, CompiledSchemaTestNode],
+    capture_macros: bool = False,
+) -> None:
+    """Render each of the test kwargs in the given context using the native
+    renderer, then insert that value into the given context as the special test
+    keyword arguments member.
+    """
+    looks_like_func = r'^\s*(env_var|ref|var|source|doc)\s*\(.+\)\s*$'
+
+    # we never care about the keypath
+    def _convert_function(value: Any, _: Any) -> Any:
+        if isinstance(value, str):
+            if re.match(looks_like_func, value) is not None:
+                # curly braces to make rendering happy
+                value = f'{{{{ {value} }}}}'
+
+            value = get_rendered(
+                value, context, node, capture_macros=capture_macros,
+                native=True
+            )
+
+        return value
+
+    kwargs = deep_map(_convert_function, node.test_metadata.kwargs)
+    context[SCHEMA_TEST_KWARGS_NAME] = kwargs

core/dbt/compilation.py

@@ -1,7 +1,7 @@
 import itertools
 import os
 from collections import defaultdict
-from typing import List, Dict
+from typing import List, Dict, Any
 
 import dbt.utils
 import dbt.include
@@ -11,11 +11,17 @@
 from dbt.linker import Linker
 
 from dbt.context.providers import generate_runtime_model
+from dbt.contracts.graph.manifest import Manifest
 import dbt.contracts.project
 import dbt.exceptions
 import dbt.flags
 import dbt.config
-from dbt.contracts.graph.compiled import InjectedCTE, COMPILED_TYPES
+from dbt.contracts.graph.compiled import (
+    InjectedCTE,
+    COMPILED_TYPES,
+    NonSourceCompiledNode,
+    CompiledSchemaTestNode,
+)
 from dbt.contracts.graph.parsed import ParsedNode
 
 from dbt.logger import GLOBAL_LOGGER as logger
@@ -24,12 +30,12 @@
 
 
 def _compiled_type_for(model: ParsedNode):
-    if model.resource_type not in COMPILED_TYPES:
+    if type(model) not in COMPILED_TYPES:
         raise dbt.exceptions.InternalException(
             'Asked to compile {} node, but it has no compiled form'
-            .format(model.resource_type)
+            .format(type(model))
         )
-    return COMPILED_TYPES[model.resource_type]
+    return COMPILED_TYPES[type(model)]
 
 
 def print_compile_stats(stats):
@@ -130,6 +136,22 @@ def initialize(self):
         dbt.clients.system.make_directory(self.config.target_path)
         dbt.clients.system.make_directory(self.config.modules_path)
 
+    def _create_node_context(
+        self,
+        node: NonSourceCompiledNode,
+        manifest: Manifest,
+        extra_context: Dict[str, Any],
+    ) -> Dict[str, Any]:
+        context = generate_runtime_model(
+            node, self.config, manifest
+        )
+        context.update(extra_context)
+        if isinstance(node, CompiledSchemaTestNode):
+            # for test nodes, add a special keyword args value to the context
+            dbt.clients.jinja.add_rendered_test_kwargs(context, node)
+
+        return context
+
     def compile_node(self, node, manifest, extra_context=None):
         if extra_context is None:
             extra_context = {}
@@ -146,10 +168,9 @@ def compile_node(self, node, manifest, extra_context=None):
         })
         compiled_node = _compiled_type_for(node).from_dict(data)
 
-        context = generate_runtime_model(
-            compiled_node, self.config, manifest
+        context = self._create_node_context(
+            compiled_node, manifest, extra_context
         )
-        context.update(extra_context)
 
         compiled_node.compiled_sql = dbt.clients.jinja.get_rendered(
             node.raw_sql,
@@ -160,33 +181,6 @@ def compile_node(self, node, manifest, extra_context=None):
 
         injected_node, _ = prepend_ctes(compiled_node, manifest)
 
-        should_wrap = {NodeType.Test, NodeType.Operation}
-        if injected_node.resource_type in should_wrap:
-            # data tests get wrapped in count(*)
-            # TODO : move this somewhere more reasonable
-            if 'data' in injected_node.tags and \
-               injected_node.resource_type == NodeType.Test:
-                injected_node.wrapped_sql = (
-                    "select count(*) as errors "
-                    "from (\n{test_sql}\n) sbq").format(
-                        test_sql=injected_node.injected_sql)
-            else:
-                # don't wrap schema tests or analyses.
-                injected_node.wrapped_sql = injected_node.injected_sql
-
-        elif injected_node.resource_type == NodeType.Snapshot:
-            # unfortunately we do everything automagically for
-            # snapshots. in the future it'd be nice to generate
-            # the SQL at the parser level.
-            pass
-
-        elif(injected_node.resource_type == NodeType.Model and
-             injected_node.get_materialization() == 'ephemeral'):
-            pass
-
-        else:
-            injected_node.wrapped_sql = None
-
         return injected_node
 
     def write_graph_file(self, linker, manifest):

core/dbt/config/project.py

@@ -274,7 +274,7 @@ def _preprocess(project_dict: Dict[str, Any]) -> Dict[str, Any]:
         """Pre-process certain special keys to convert them from None values
         into empty containers, and to turn strings into arrays of strings.
         """
-        handlers: Dict[Tuple[str, ...], Callable[[Any], Any]] = {
+        handlers: Dict[Tuple[Union[str, int], ...], Callable[[Any], Any]] = {
             ('on-run-start',): _list_if_none_or_string,
             ('on-run-end',): _list_if_none_or_string,
         }
@@ -286,7 +286,7 @@ def _preprocess(project_dict: Dict[str, Any]) -> Dict[str, Any]:
             handlers[(k, 'post-hook')] = _list_if_none_or_string
         handlers[('seeds', 'column_types')] = _dict_if_none
 
-        def converter(value: Any, keypath: Tuple[str, ...]) -> Any:
+        def converter(value: Any, keypath: Tuple[Union[str, int], ...]) -> Any:
             if keypath in handlers:
                 handler = handlers[keypath]
                 return handler(value)