Repository: dbt-labs/dbt-core

Issue:
RPC Server: Support macros in execute and compile tasks
## Feature

### Feature description
The RPC Server `compile/` and `execute/` endpoints should accept macros defined inline in raw SQL. Our current strategy of sending the same text to the `macros` and `sql` arguments of the endpoints doesn't work because of the macro cache. The following query results in `maximum recursion depth exceeded`:

```
{{ test_macros() }}

{% macro test_macros() %}
  1
{% endmacro %}
```

The reason this happens is because we re-parse the entire contents of the macro file at runtime, so we end up in an infinite recursion scenario.


Discussion:

Resolving patch:
diff --git a/core/dbt/clients/_jinja_blocks.py b/core/dbt/clients/_jinja_blocks.py
new file mode 100644
--- /dev/null
+++ b/core/dbt/clients/_jinja_blocks.py
@@ -0,0 +1,380 @@
+import re
+
+import dbt.exceptions
+
+
+def regex(pat):
+    return re.compile(pat, re.DOTALL | re.MULTILINE)
+
+
+class BlockData(object):
+    """raw plaintext data from the top level of the file."""
+    def __init__(self, contents):
+        self.block_type_name = '__dbt__data'
+        self.contents = contents
+        self.full_block = contents
+
+
+class BlockTag(object):
+    def __init__(self, block_type_name, block_name, contents=None, **kw):
+        self.block_type_name = block_type_name
+        self.block_name = block_name
+        self.contents = contents
+        self.full_block = None
+
+    def __str__(self):
+        return 'BlockTag({!r}, {!r})'.format(self.block_type_name,
+                                             self.block_name)
+
+    def __repr__(self):
+        return str(self)
+
+    @property
+    def end_block_type_name(self):
+        return 'end{}'.format(self.block_type_name)
+
+    def end_pat(self):
+        # we don't want to use string formatting here because jinja uses most
+        # of the string formatting operators in its syntax...
+        pattern = ''.join((
+            r'(?P<endblock>((?:\s*\{\%\-|\{\%)\s*',
+            self.end_block_type_name,
+            r'\s*(?:\-\%\}\s*|\%\})))',
+        ))
+        return regex(pattern)
+
+
+_NAME_PATTERN = r'[A-Za-z_][A-Za-z_0-9]*'
+
+COMMENT_START_PATTERN = regex(r'(?:(?P<comment_start>(\s*\{\#)))')
+COMMENT_END_PATTERN = regex(r'(.*?)(\s*\#\})')
+RAW_START_PATTERN = regex(
+    r'(?:\s*\{\%\-|\{\%)\s*(?P<raw_start>(raw))\s*(?:\-\%\}\s*|\%\})'
+)
+
+BLOCK_START_PATTERN = regex(''.join((
+    r'(?:\s*\{\%\-|\{\%)\s*',
+    r'(?P<block_type_name>({}))'.format(_NAME_PATTERN),
+    # some blocks have a 'block name'.
+    r'(?:\s+(?P<block_name>({})))?'.format(_NAME_PATTERN),
+)))
+
+TAG_CLOSE_PATTERN = regex(r'(?:\-\%\}\s*|\%\})')
+# if you do {% materialization foo, adapter="myadapter' %} and end up with
+# mismatched quotes this will still match, but jinja will fail somewhere
+# since the adapter= argument has to be an adapter name, and none have quotes
+# or anything else in them. So this should be fine.
+MATERIALIZATION_ARGS_PATTERN = regex(
+    r'\s*,\s*'
+    r'''(?P<adpater_arg>(adapter=(?:['"]{}['"])|default))'''
+    .format(_NAME_PATTERN)
+)
+# macros an stuff like macros get open parents, followed by a very complicated
+# argument spec! In fact, it's easiest to parse it in tiny little chunks
+# because we have to handle awful stuff like string parsing ;_;
+MACRO_ARGS_START_PATTERN = regex(r'\s*(?P<macro_start>\()\s*')
+MACRO_ARGS_END_PATTERN = regex(r'\s*(?P<macro_end>(\)))\s*')
+
+# macros can be like {% macro foo(bar) %} or {% macro foo(bar, baz) %} or
+# {% macro foo(bar, baz="quux") %} or ...
+# I think jinja disallows default values after required (like Python), but we
+# can ignore that and let jinja deal
+MACRO_ARG_PATTERN = regex(''.join((
+    r'\s*(?P<macro_arg_name>({}))\s*',
+    r'((?P<value>=)|(?P<more_args>,)?)\s*'.format(_NAME_PATTERN),
+)))
+
+# stolen from jinja's lexer. Note that we've consumed all prefix whitespace by
+# the time we want to use this.
+STRING_PATTERN = regex(
+    r"(?P<string>('([^'\\]*(?:\\.[^'\\]*)*)'|"
+    r'"([^"\\]*(?:\\.[^"\\]*)*)"))'
+)
+
+# any number of non-quote characters, followed by:
+# - quote: a quote mark indicating start of a string (you'll want to backtrack
+#          the regex end on quotes and then match with the string pattern)
+# - a comma (so there will be another full argument)
+# - a closing parenthesis (you can now expect a closing tag)
+NON_STRING_MACRO_ARGS_PATTERN = regex(
+    # anything, followed by a quote, open/close paren, or comma
+    r'''(.*?)'''
+    r'''((?P<quote>(['"]))|(?P<open>(\())|(?P<close>(\)))|(?P<comma>(\,)))'''
+)
+
+
+class BlockIterator(object):
+    def __init__(self, data):
+        self.data = data
+        self.blocks = []
+        self._block_contents = None
+        self._parenthesis_stack = []
+        self.pos = 0
+
+    def advance(self, new_position):
+        blk = self.data[self.pos:new_position]
+
+        if self._block_contents is not None:
+            self._block_contents += blk
+
+        self.pos = new_position
+
+    def rewind(self, amount=1):
+        if self._block_contents is not None:
+            self._block_contents = self._block_contents[:-amount]
+
+        self.pos -= amount
+
+    def _search(self, pattern):
+        return pattern.search(self.data, self.pos)
+
+    def _match(self, pattern):
+        return pattern.match(self.data, self.pos)
+
+    def expect_comment_end(self):
+        """Expect a comment end and return the match object.
+        """
+        match = self._match(COMMENT_END_PATTERN)
+        if match is None:
+            dbt.exceptions.raise_compiler_error('unexpected EOF, expected #}')
+        self.advance(match.end())
+
+    def expect_raw_end(self):
+        end_pat = BlockTag('raw', None).end_pat()
+        match = self._search(end_pat)
+        if match is None:
+            dbt.exceptions.raise_compiler_error(
+                'unexpected EOF, expected {% endraw %}'
+            )
+        self.advance(match.end())
+
+    def _first_match(self, *patterns, **kwargs):
+        matches = []
+        for pattern in patterns:
+            # default to 'search', but sometimes we want to 'match'.
+            if kwargs.get('method', 'search') == 'search':
+                match = self._search(pattern)
+            else:
+                match = self._match(pattern)
+            if match:
+                matches.append(match)
+        if not matches:
+            return None
+        # if there are multiple matches, pick the least greedy match
+        # TODO: do I need to account for m.start(), or is this ok?
+        return min(matches, key=lambda m: m.end())
+
+    def _expect_match(self, expected_name, *patterns, **kwargs):
+        match = self._first_match(*patterns, **kwargs)
+        if match is None:
+            msg = 'unexpected EOF, expected {}, got "{}"'.format(
+                    expected_name, self.data[self.pos:]
+                )
+            dbt.exceptions.raise_compiler_error(msg)
+        return match
+
+    def handle_block(self, match, block_start=None):
+        """Handle a block. The current state of the parser should be after the
+        open block is completed:
+            {% blk foo %}my data {% endblk %}
+                         ^ right here
+        """
+        # we have to handle comments inside blocks because you could do this:
+        # {% blk foo %}asdf {# {% endblk %} #} {%endblk%}
+        # they still end up in the data/raw_data of the block itself, but we
+        # have to know to ignore stuff until the end comment marker!
+        found = BlockTag(**match.groupdict())
+        # the full block started at the given match start, which may include
+        # prefixed whitespace! we'll strip it later
+        if block_start is None:
+            block_start = match.start()
+
+        self._block_contents = ''
+
+        # you can have as many comments in your block as you'd like!
+        while True:
+            match = self._expect_match(
+                '"{}"'.format(found.end_block_type_name),
+                found.end_pat(), COMMENT_START_PATTERN, RAW_START_PATTERN
+            )
+            groups = match.groupdict()
+            if groups.get('endblock') is not None:
+                break
+
+            self.advance(match.end())
+
+            if groups.get('comment_start') is not None:
+                self.expect_comment_end()
+            elif groups.get('raw_start') is not None:
+                self.expect_raw_end()
+            else:
+                raise dbt.exceptions.InternalException(
+                    'unhandled regex in handle_block, no match: {}'
+                    .format(groups)
+                )
+
+        # we want to advance to just the end tag at first, to extract the
+        # contents
+        self.advance(match.start())
+        found.contents = self._block_contents
+        self._block_contents = None
+        # now advance to the end
+        self.advance(match.end())
+        found.full_block = self.data[block_start:self.pos]
+        return found
+
+    def handle_materialization(self, match):
+        self._expect_match('materialization args',
+                           MATERIALIZATION_ARGS_PATTERN)
+        self._expect_match('%}', TAG_CLOSE_PATTERN)
+        # handle the block we started with!
+        self.blocks.append(self.handle_block(match))
+
+    def find_block(self):
+        open_block = (
+            r'(?:\s*\{\%\-|\{\%)\s*'
+            r'(?P<block_type_name>([A-Za-z_][A-Za-z_0-9]*))'
+            # some blocks have a 'block name'.
+            r'(?:\s+(?P<block_name>([A-Za-z_][A-Za-z_0-9]*)))?'
+        )
+
+        match = self._first_match(regex(open_block), COMMENT_START_PATTERN)
+        if match is None:
+            return False
+
+        raw_toplevel = self.data[self.pos:match.start()]
+        if len(raw_toplevel) > 0:
+            self.blocks.append(BlockData(raw_toplevel))
+
+        matchgroups = match.groupdict()
+
+        # comments are easy
+        if matchgroups.get('comment_start') is not None:
+            start = match.start()
+            self.expect_comment_end()
+            self.blocks.append(BlockData(self.data[start:self.pos]))
+            return True
+
+        if matchgroups.get('block_type_name') == 'raw':
+            start = match.start()
+            self.expect_raw_end()
+            self.blocks.append(BlockData(self.data[start:self.pos]))
+            return True
+
+        if matchgroups.get('block_type_name') == 'materialization':
+            self.advance(match.end())
+            self.handle_materialization(match)
+            return True
+
+        # we're somewhere like this {% block_type_name block_type
+        # we've either got arguments, a close of tag (%}), or bad input.
+        # we've handled materializations already (they're weird!)
+        # thankfully, comments aren't allowed *inside* a block def...
+        block_end_match = self._expect_match('%} or (...)',
+                                             TAG_CLOSE_PATTERN,
+                                             MACRO_ARGS_START_PATTERN)
+        self.advance(block_end_match.end())
+        if block_end_match.groupdict().get('macro_start') is not None:
+            # we've hit our first parenthesis!
+            self._parenthesis_stack = [True]
+            self._process_macro_args()
+            self.advance(self._expect_match('%}', TAG_CLOSE_PATTERN).end())
+
+        # tag close time!
+        self.blocks.append(self.handle_block(match))
+        return True
+
+    def _process_macro_default_arg(self):
+        """Handle the bit after an '=' in a macro default argument. This is
+        probably the trickiest thing. The goal here is to accept all strings
+        jinja would accept and always handle block start/end correctly: It's
+        fine to have false positives, jinja can fail later.
+
+        Return True if there are more arguments expected.
+        """
+        while self._parenthesis_stack:
+            match = self._expect_match(
+                'macro argument',
+                # you could have a string
+                STRING_PATTERN,
+                # a quote, a comma, or a close parenthesis
+                NON_STRING_MACRO_ARGS_PATTERN,
+                # we want to "match", not "search"
+                method='match'
+            )
+            matchgroups = match.groupdict()
+            self.advance(match.end())
+            if matchgroups.get('string') is not None:
+                # we got a string value. There could be more data.
+                continue
+            elif matchgroups.get('quote') is not None:
+                # we got a bunch of data and then a string opening value.
+                # put the quote back on the menu
+                self.rewind()
+                # if we only got a single quote mark and didn't hit a string
+                # at all, this file has an unclosed quote. Fail accordingly.
+                if match.end() - match.start() == 1:
+                    msg = (
+                        'Unclosed quotation mark at position {}. Context:\n{}'
+                        .format(self.pos, self.data[self.pos-20:self.pos+20])
+                    )
+                    dbt.exceptions.raise_compiler_error(msg)
+            elif matchgroups.get('comma') is not None:
+                # small hack: if we hit a comma and there is one parenthesis
+                # left, return to look for a new name. otherwise we're still
+                # looking for the parameter close.
+                if len(self._parenthesis_stack) == 1:
+                    return
+            elif matchgroups.get('close'):
+                self._parenthesis_stack.pop()
+            else:
+                raise dbt.exceptions.InternalException(
+                    'unhandled regex in _process_macro_default_arg(), no match'
+                    ': {}'.format(matchgroups)
+                )
+
+    def _process_macro_args(self):
+        """Macro args are pretty tricky! Arg names themselves are simple, but
+        you can set arbitrary default values, including doing stuff like:
+        {% macro my_macro(arg="x" + ("}% {# {% endmacro %}" * 2)) %}
+
+        Which makes you a jerk, but is valid jinja.
+        """
+        # we are currently after the first parenthesis (+ any whitespace) after
+        # the macro args started. You can either have the close paren, or a
+        # name.
+        while self._parenthesis_stack:
+            match = self._expect_match('macro arguments',
+                                       MACRO_ARGS_END_PATTERN,
+                                       MACRO_ARG_PATTERN)
+            self.advance(match.end())
+            matchgroups = match.groupdict()
+            if matchgroups.get('macro_end') is not None:
+                self._parenthesis_stack.pop()
+            # we got an argument. let's see what it has
+            elif matchgroups.get('value') is not None:
+                # we have to process a single macro argument. This mutates
+                # the parenthesis stack! If it finds a comma, it will continue
+                # the loop.
+                self._process_macro_default_arg()
+            elif matchgroups.get('more_args') is not None:
+                continue
+            else:
+                raise dbt.exceptions.InternalException(
+                    'unhandled regex in _process_macro_args(), no match: {}'
+                    .format(matchgroups)
+                )
+            # if there are more arguments or a macro arg end we'll catch them
+            # on the next loop around
+
+    def lex_for_blocks(self):
+        while self.data[self.pos:]:
+            found = self.find_block()
+            if not found:
+                break
+
+        raw_toplevel = self.data[self.pos:]
+        if len(raw_toplevel) > 0:
+            self.blocks.append(BlockData(raw_toplevel))
+
+        return self.blocks
diff --git a/core/dbt/clients/jinja.py b/core/dbt/clients/jinja.py
--- a/core/dbt/clients/jinja.py
+++ b/core/dbt/clients/jinja.py
@@ -12,8 +12,7 @@
 import dbt.compat
 import dbt.exceptions
 
-from dbt.node_types import NodeType
-from dbt.utils import AttrDict
+from dbt.clients._jinja_blocks import BlockIterator
 
 from dbt.logger import GLOBAL_LOGGER as logger  # noqa
 
@@ -281,3 +280,7 @@ def get_rendered(string, ctx, node=None,
 
 def undefined_error(msg):
     raise jinja2.exceptions.UndefinedError(msg)
+
+
+def extract_toplevel_blocks(data):
+    return BlockIterator(data).lex_for_blocks()
diff --git a/core/dbt/task/compile.py b/core/dbt/task/compile.py
--- a/core/dbt/task/compile.py
+++ b/core/dbt/task/compile.py
@@ -1,6 +1,7 @@
 import os
 
 from dbt.adapters.factory import get_adapter
+from dbt.clients.jinja import extract_toplevel_blocks
 from dbt.compilation import compile_manifest
 from dbt.loader import load_all_projects, GraphLoader
 from dbt.node_runners import CompileRunner, RPCCompileRunner
@@ -56,12 +57,26 @@ def runtime_cleanup(self, selected_uids):
         self._skipped_children = {}
         self._raise_next_tick = None
 
-    def handle_request(self, name, sql, macros=None):
+    def _extract_request_data(self, data):
+        data = self.decode_sql(data)
+        macro_blocks = []
+        data_chunks = []
+        for block in extract_toplevel_blocks(data):
+            if block.block_type_name == 'macro':
+                macro_blocks.append(block.full_block)
+            else:
+                data_chunks.append(block.full_block)
+        macros = '\n'.join(macro_blocks)
+        sql = ''.join(data_chunks)
+        return sql, macros
+
+    def handle_request(self, name, sql):
         request_path = os.path.join(self.config.target_path, 'rpc', name)
         all_projects = load_all_projects(self.config)
         macro_overrides = {}
-        if macros is not None:
-            macros = self.decode_sql(macros)
+        sql, macros = self._extract_request_data(sql)
+
+        if macros:
             macro_parser = MacroParser(self.config, all_projects)
             macro_overrides.update(macro_parser.parse_macro_file(
                 macro_file_path='from remote system',
@@ -77,7 +92,6 @@ def handle_request(self, name, sql, macros=None):
             macro_manifest=self._base_manifest
         )
 
-        sql = self.decode_sql(sql)
         node_dict = {
             'name': name,
             'root_path': request_path,


Test patch:
diff --git a/test/integration/042_sources_test/test_sources.py b/test/integration/042_sources_test/test_sources.py
--- a/test/integration/042_sources_test/test_sources.py
+++ b/test/integration/042_sources_test/test_sources.py
@@ -322,11 +322,16 @@ def project_config(self):
         }
 
     def build_query(self, method, kwargs, sql=None, test_request_id=1, macros=None):
+        body_data = ''
         if sql is not None:
-            kwargs['sql'] = b64(sql.encode('utf-8')).decode('utf-8')
+            body_data += sql
 
         if macros is not None:
-            kwargs['macros'] = b64(macros.encode('utf-8')).decode('utf-8')
+            body_data += macros
+
+        if sql is not None or macros is not None:
+            kwargs['sql'] = b64(body_data.encode('utf-8')).decode('utf-8')
+
         return {
             'jsonrpc': '2.0',
             'method': method,
@@ -464,6 +469,18 @@ def test_compile(self):
             compiled_sql='select 2 as id'
         )
 
+        macro_override_with_if_statement = self.query(
+            'compile',
+            '{% if True %}select {{ happy_little_macro() }}{% endif %}',
+            name='foo',
+            macros='{% macro override_me() %}2 as id{% endmacro %}'
+        ).json()
+        self.assertSuccessfulCompilationResult(
+            macro_override_with_if_statement,
+            '{% if True %}select {{ happy_little_macro() }}{% endif %}',
+            compiled_sql='select 2 as id'
+        )
+
     @use_profile('postgres')
     def test_run(self):
         # seed + run dbt to make models before using them!
@@ -538,6 +555,43 @@ def test_run(self):
             table={'column_names': ['id'], 'rows': [[2.0]]}
         )
 
+        macro_override_with_if_statement = self.query(
+            'run',
+            '{% if True %}select {{ happy_little_macro() }}{% endif %}',
+            name='foo',
+            macros='{% macro override_me() %}2 as id{% endmacro %}'
+        ).json()
+        self.assertSuccessfulRunResult(
+            macro_override_with_if_statement,
+            '{% if True %}select {{ happy_little_macro() }}{% endif %}',
+            compiled_sql='select 2 as id',
+            table={'column_names': ['id'], 'rows': [[2.0]]}
+        )
+
+        macro_with_raw_statement = self.query(
+            'run',
+            '{% raw %}select 1 as{% endraw %}{{ test_macros() }}{% macro test_macros() %} id{% endmacro %}',
+            name='foo'
+        ).json()
+        self.assertSuccessfulRunResult(
+            macro_with_raw_statement,
+            '{% raw %}select 1 as{% endraw %}{{ test_macros() }}',
+            compiled_sql='select 1 as id',
+            table={'column_names': ['id'], 'rows': [[1.0]]}
+        )
+
+        macro_with_comment = self.query(
+            'run',
+            '{% raw %}select 1 {% endraw %}{{ test_macros() }} {# my comment #}{% macro test_macros() -%} as{% endmacro %} id{# another comment #}',
+            name='foo'
+        ).json()
+        self.assertSuccessfulRunResult(
+            macro_with_comment,
+            '{% raw %}select 1 {% endraw %}{{ test_macros() }} {# my comment #} id{# another comment #}',
+            compiled_sql='select 1 as  id',
+            table={'column_names': ['id'], 'rows': [[1.0]]}
+        )
+
     @use_profile('postgres')
     def test_invalid_requests(self):
         data = self.query(
diff --git a/test/unit/test_jinja.py b/test/unit/test_jinja.py
--- a/test/unit/test_jinja.py
+++ b/test/unit/test_jinja.py
@@ -1,6 +1,9 @@
 import unittest
 
 from dbt.clients.jinja import get_template
+from dbt.clients.jinja import extract_toplevel_blocks
+from dbt.exceptions import CompilationException
+
 
 class TestJinja(unittest.TestCase):
     def test_do(self):
@@ -9,3 +12,179 @@ def test_do(self):
         template = get_template(s, {})
         mod = template.make_module()
         self.assertEqual(mod.my_dict, {'a': 1})
+
+
+class TestBlockLexer(unittest.TestCase):
+    def test_basic(self):
+        body = '{{ config(foo="bar") }}\r\nselect * from this.that\r\n'
+        block_data = '  \n\r\t{%- mytype foo %}'+body+'{%endmytype -%}'
+        blocks = extract_toplevel_blocks(block_data)
+        self.assertEqual(len(blocks), 1)
+        self.assertEqual(blocks[0].block_type_name, 'mytype')
+        self.assertEqual(blocks[0].block_name, 'foo')
+        self.assertEqual(blocks[0].contents, body)
+        self.assertEqual(blocks[0].full_block, block_data)
+
+    def test_multiple(self):
+        body_one = '{{ config(foo="bar") }}\r\nselect * from this.that\r\n'
+        body_two = (
+            '{{ config(bar=1)}}\r\nselect * from {% if foo %} thing '
+            '{% else %} other_thing {% endif %}'
+        )
+
+        block_data = (
+            '  {% mytype foo %}' + body_one + '{% endmytype %}' +
+            '\r\n{% othertype bar %}' + body_two + '{% endothertype %}'
+        )
+        all_blocks = extract_toplevel_blocks(block_data)
+        blocks = [b for b in all_blocks if b.block_type_name != '__dbt__data']
+        self.assertEqual(len(blocks), 2)
+
+    def test_comments(self):
+        body = '{{ config(foo="bar") }}\r\nselect * from this.that\r\n'
+        comment = '{# my comment #}'
+        block_data = '  \n\r\t{%- mytype foo %}'+body+'{%endmytype -%}'
+        all_blocks = extract_toplevel_blocks(comment+block_data)
+        blocks = [b for b in all_blocks if b.block_type_name != '__dbt__data']
+        self.assertEqual(len(blocks), 1)
+        self.assertEqual(blocks[0].block_type_name, 'mytype')
+        self.assertEqual(blocks[0].block_name, 'foo')
+        self.assertEqual(blocks[0].contents, body)
+        self.assertEqual(blocks[0].full_block, block_data)
+
+    def test_evil_comments(self):
+        body = '{{ config(foo="bar") }}\r\nselect * from this.that\r\n'
+        comment = '{# external comment {% othertype bar %} select * from thing.other_thing{% endothertype %} #}'
+        block_data = '  \n\r\t{%- mytype foo %}'+body+'{%endmytype -%}'
+        all_blocks = extract_toplevel_blocks(comment+block_data)
+        blocks = [b for b in all_blocks if b.block_type_name != '__dbt__data']
+        self.assertEqual(len(blocks), 1)
+        self.assertEqual(blocks[0].block_type_name, 'mytype')
+        self.assertEqual(blocks[0].block_name, 'foo')
+        self.assertEqual(blocks[0].contents, body)
+        self.assertEqual(blocks[0].full_block, block_data)
+
+    def test_nested_comments(self):
+        body = '{# my comment #} {{ config(foo="bar") }}\r\nselect * from {# my other comment embedding {% endmytype %} #} this.that\r\n'
+        block_data = '  \n\r\t{%- mytype foo %}'+body+'{% endmytype -%}'
+        comment = '{# external comment {% othertype bar %} select * from thing.other_thing{% endothertype %} #}'
+        all_blocks = extract_toplevel_blocks(comment+block_data)
+        blocks = [b for b in all_blocks if b.block_type_name != '__dbt__data']
+        self.assertEqual(len(blocks), 1)
+        self.assertEqual(blocks[0].block_type_name, 'mytype')
+        self.assertEqual(blocks[0].block_name, 'foo')
+        self.assertEqual(blocks[0].contents, body)
+        self.assertEqual(blocks[0].full_block, block_data)
+
+    def test_complex_file(self):
+        all_blocks = extract_toplevel_blocks(complex_archive_file)
+        blocks = [b for b in all_blocks if b.block_type_name != '__dbt__data']
+        self.assertEqual(len(blocks), 3)
+        self.assertEqual(blocks[0].block_type_name, 'mytype')
+        self.assertEqual(blocks[0].block_name, 'foo')
+        self.assertEqual(blocks[0].full_block, '{% mytype foo %} some stuff {% endmytype %}')
+        self.assertEqual(blocks[0].contents, ' some stuff ')
+        self.assertEqual(blocks[1].block_type_name, 'mytype')
+        self.assertEqual(blocks[1].block_name, 'bar')
+        self.assertEqual(blocks[1].full_block, bar_block)
+        self.assertEqual(blocks[1].contents, bar_block[16:-15].rstrip())
+        self.assertEqual(blocks[2].block_type_name, 'myothertype')
+        self.assertEqual(blocks[2].block_name, 'x')
+        self.assertEqual(blocks[2].full_block, x_block.strip())
+        self.assertEqual(blocks[2].contents, x_block[len('\n{% myothertype x %}'):-len('{% endmyothertype %}\n')])
+
+    def test_peaceful_macro_coexistence(self):
+        body = '{# my macro #} {% macro foo(a, b) %} do a thing {%- endmacro %} {# my model #} {% a b %} {% enda %}'
+        all_blocks = extract_toplevel_blocks(body)
+        blocks = [b for b in all_blocks if b.block_type_name != '__dbt__data']
+        self.assertEqual(len(blocks), 2)
+        self.assertEqual(blocks[0].block_type_name, 'macro')
+        self.assertEqual(blocks[0].block_name, 'foo')
+        self.assertEqual(blocks[0].contents, ' do a thing')
+        self.assertEqual(blocks[1].block_type_name, 'a')
+        self.assertEqual(blocks[1].block_name, 'b')
+        self.assertEqual(blocks[1].contents, ' ')
+
+    def test_macro_with_crazy_args(self):
+        body = '''{% macro foo(a, b=asdf("cool this is 'embedded'" * 3) + external_var, c)%}cool{# block comment with {% endmacro %} in it #} stuff here {% endmacro %}'''
+        all_blocks = extract_toplevel_blocks(body)
+        blocks = [b for b in all_blocks if b.block_type_name != '__dbt__data']
+        self.assertEqual(len(blocks), 1)
+        self.assertEqual(blocks[0].block_type_name, 'macro')
+        self.assertEqual(blocks[0].block_name, 'foo')
+        self.assertEqual(blocks[0].contents, 'cool{# block comment with {% endmacro %} in it #} stuff here ')
+
+    def test_materialization_parse(self):
+        body = '{% materialization xxx, default %} ... {% endmaterialization %}'
+        all_blocks = extract_toplevel_blocks(body)
+        blocks = [b for b in all_blocks if b.block_type_name != '__dbt__data']
+        self.assertEqual(len(blocks), 1)
+        self.assertEqual(blocks[0].block_type_name, 'materialization')
+        self.assertEqual(blocks[0].block_name, 'xxx')
+        self.assertEqual(blocks[0].full_block, body)
+
+        body = '{% materialization xxx, adapter="other" %} ... {% endmaterialization %}'
+        all_blocks = extract_toplevel_blocks(body)
+        blocks = [b for b in all_blocks if b.block_type_name != '__dbt__data']
+        self.assertEqual(len(blocks), 1)
+        self.assertEqual(blocks[0].block_type_name, 'materialization')
+        self.assertEqual(blocks[0].block_name, 'xxx')
+        self.assertEqual(blocks[0].full_block, body)
+
+    def test_nested_failure(self):
+        # we don't allow nesting same blocks
+        # ideally we would not allow nesting any, but that's much harder
+        body = '{% myblock a %} {% myblock b %} {% endmyblock %} {% endmyblock %}'
+        with self.assertRaises(CompilationException):
+            extract_toplevel_blocks(body)
+
+    def test_incomplete_block_failure(self):
+        fullbody = '{% myblock foo %} {% endblock %}'
+        for length in range(1, len(fullbody)-1):
+            body = fullbody[:length]
+        with self.assertRaises(CompilationException):
+            extract_toplevel_blocks(body)
+
+    def test_wrong_end_failure(self):
+        body = '{% myblock foo %} {% endotherblock %}'
+        with self.assertRaises(CompilationException):
+            extract_toplevel_blocks(body)
+
+    def test_comment_no_end_failure(self):
+        body = '{# '
+        with self.assertRaises(CompilationException):
+            extract_toplevel_blocks(body)
+
+    def test_comment_only(self):
+        body = '{# myblock #}'
+        all_blocks = extract_toplevel_blocks(body)
+        blocks = [b for b in all_blocks if b.block_type_name != '__dbt__data']
+        self.assertEqual(len(blocks), 0)
+
+bar_block = '''{% mytype bar %}
+{# a comment
+    that inside it has
+    {% mytype baz %}
+{% endmyothertype %}
+{% endmytype %}
+{% endmytype %}
+    {#
+{% endmytype %}#}
+
+some other stuff
+
+{%- endmytype%}'''
+
+x_block = '''
+{% myothertype x %}
+before
+{##}
+and after
+{% endmyothertype %}
+'''
+
+complex_archive_file = '''
+{#some stuff {% mytype foo %} #}
+{% mytype foo %} some stuff {% endmytype %}
+
+'''+bar_block+x_block