Feat!: Extend SQL parser with designated blocks for jinja code (#963)

Author: izeigerman

docs/concepts/macros/jinja_macros.md

@@ -1,8 +1,8 @@
 # Jinja macros
 
-SQLMesh supports macros from the [Jinja](https://jinja.palletsprojects.com/en/3.1.x/) templating system. 
+SQLMesh supports macros from the [Jinja](https://jinja.palletsprojects.com/en/3.1.x/) templating system.
 
-Jinja's macro approach is pure string substitution. Unlike SQLMesh macros, they assemble SQL query text without building a semantic representation. 
+Jinja's macro approach is pure string substitution. Unlike SQLMesh macros, they assemble SQL query text without building a semantic representation.
 
 **NOTE:** SQLMesh projects support the standard Jinja function library only - they do **not** support dbt-specific jinja functions like `{{ ref() }}`. dbt-specific functions are allowed in dbt projects being run with the [SQLMesh adapter](../../integrations/dbt.md).
 
@@ -16,6 +16,41 @@ The three curly brace symbols are:
 - `{%...%}` creates Jinja statements. Statements give instructions to Jinja, such as setting variable values, control flow with `if`, `for` loops, and defining macro functions.
 - `{#...#}` creates Jinja comments. These comments will not be included in the rendered SQL query.
 
+Since Jinja strings are not syntactically valid SQL expressions and cannot be parsed as such, the model query must be wrapped in a special `JINJA_QUERY_BEGIN; ...; JINJA_END;` block in order for SQLMesh to detect it:
+
+```sql linenums="1" hl_lines="5 9"
+MODEL (
+  name sqlmesh_example.full_model
+);
+
+JINJA_QUERY_BEGIN;
+
+SELECT {{ 1 + 1 }};
+
+JINJA_END;
+```
+
+Similarly, to use Jinja expressions as part of statements that should be evaluated before or after the model query, the `JINJA_STATEMENT_BEGIN; ...; JINJA_END;` block should be used:
+
+```sql linenums="1"
+MODEL (
+  name sqlmesh_example.full_model
+);
+
+JINJA_STATEMENT_BEGIN;
+{{ pre_hook() }}
+JINJA_END;
+
+JINJA_QUERY_BEGIN;
+SELECT {{ 1 + 1 }};
+JINJA_END;
+
+JINJA_STATEMENT_BEGIN;
+{{ post_hook() }}
+JINJA_END;
+```
+
+
 ## User-defined variables
 
 Define your own variables with the Jinja statement `{% set ... %}`. For example, we could specify the name of the `num_orders` column in the `sqlmesh_example.full_model` like this:
@@ -28,6 +63,8 @@ MODEL (
   audits [assert_positive_order_ids],
 );
 
+JINJA_QUERY_BEGIN;
+
 {% set my_col = 'num_orders' %} -- Jinja definition of variable `my_col`
 
 SELECT
@@ -36,6 +73,8 @@ SELECT
 FROM
     sqlmesh_example.incremental_model
 GROUP BY item_id
+
+JINJA_END;
 ```
 
 Note that the Jinja set statement is written after the `MODEL` statement and before the SQL query.
@@ -48,7 +87,7 @@ Jinja variables can be string, integer, or float data types. They can also be an
 
 #### for loops
 
-For loops let you iterate over a collection of items to condense repetitive code and easily change the values used by the code. 
+For loops let you iterate over a collection of items to condense repetitive code and easily change the values used by the code.
 
 Jinja for loops begin with `{% for ... %}` and end with `{% endfor %}`. This example demonstrates creating indicator variables with `CASE WHEN` using a Jinja for loop:
 
@@ -88,9 +127,9 @@ FROM table
 
 The rendered query would be the same as before.
 
-#### if 
+#### if
 
-if statements allow you to take an action (or not) based on some condition. 
+if statements allow you to take an action (or not) based on some condition.
 
 Jinja if statements begin with `{% if ... %}` and end with `{% endif %}`. The starting `if` statement must contain code that evaluates to `True` or `False`. For example, all of `True`, `1 + 1 == 2`, and `'a' in ['a', 'b']` evaluate to `True`.
 
@@ -118,11 +157,11 @@ FROM table
 
 ## User-defined macro functions
 
-User-defined macro functions allow the same macro code to be used in multiple models. 
+User-defined macro functions allow the same macro code to be used in multiple models.
 
 Jinja macro functions should be placed in `.sql` files in the SQLMesh project's `macros` directory. Multiple functions can be defined in one `.sql` file, or they can be distributed across multiple files.
 
-Jinja macro functions are defined with the `{% macro %}` and `{% endmacro %}` statements. The macro function name and arguments are specified in the `{% macro %}` statement. 
+Jinja macro functions are defined with the `{% macro %}` and `{% endmacro %}` statements. The macro function name and arguments are specified in the `{% macro %}` statement.
 
 For example, a macro function named `print_text` that takes no arguments could be defined with:
 
@@ -186,6 +225,6 @@ Some SQL dialects interpret double and single quotes differently. We could repla
 
 ## Mixing macro systems
 
-SQLMesh supports both the Jinja and [SQLMesh](./sqlmesh_macros.md) macro systems. We strongly recommend using only one system in a single model - if both are present, they may fail or behave in unintuitive ways. 
+SQLMesh supports both the Jinja and [SQLMesh](./sqlmesh_macros.md) macro systems. We strongly recommend using only one system in a single model - if both are present, they may fail or behave in unintuitive ways.
 
 [Predefined SQLMesh macro variables](./macro_variables.md) can be used in a query containing user-defined Jinja variables and functions. However, predefined variables passed as arguments to a user-defined Jinja macro function must use the Jinja curly brace syntax `{{ start_ds }}` instead of the SQLMesh macro `@` prefix syntax `@start_ds`. Note that curly brace syntax may require quoting to generate the equivalent of the `@` syntax.

examples/sushi/models/waiter_as_customer_by_day.sql

@@ -12,6 +12,8 @@ MODEL (
   )
 );
 
+JINJA_QUERY_BEGIN;
+
 {% set x = 1 %}
 
 SELECT
@@ -21,4 +23,6 @@ SELECT
   {{ alias(identity(x), 'flag') }}
 FROM sushi.waiters AS w
 JOIN sushi.customers as c ON w.waiter_id = c.customer_id
-JOIN sushi.waiter_names as wn ON w.waiter_id = wn.id
+JOIN sushi.waiter_names as wn ON w.waiter_id = wn.id;
+
+JINJA_END;

sqlmesh/core/audit/definition.py

@@ -53,7 +53,7 @@ class Audit(AuditMeta, frozen=True):
     An audit is a SQL query that returns bad records.
     """
 
-    query: t.Union[exp.Subqueryable, d.Jinja]
+    query: t.Union[exp.Subqueryable, d.JinjaQuery]
     expressions_: t.Optional[t.List[exp.Expression]] = Field(default=None, alias="expressions")
     jinja_macros: JinjaMacroRegistry = JinjaMacroRegistry()
 

sqlmesh/core/dialect.py

@@ -4,12 +4,14 @@
 import re
 import typing as t
 from difflib import unified_diff
+from enum import Enum, auto
 
 import pandas as pd
 from jinja2.meta import find_undeclared_variables
 from sqlglot import Dialect, Generator, Parser, TokenType, exp
 from sqlglot.dialects.dialect import DialectType
 from sqlglot.optimizer.normalize_identifiers import normalize_identifiers
+from sqlglot.tokens import Token
 
 from sqlmesh.core.constants import MAX_MODEL_DEFINITION_SIZE
 from sqlmesh.utils.jinja import ENVIRONMENT
@@ -29,6 +31,14 @@ class Jinja(exp.Func):
     is_var_len_args = True
 
 
+class JinjaQuery(Jinja):
+    pass
+
+
+class JinjaStatement(Jinja):
+    pass
+
+
 class ModelKind(exp.Expression):
     arg_types = {"this": True, "expressions": False}
 
@@ -387,6 +397,47 @@ def text_diff(
 )
 
 
+def _is_command_statement(command: str, tokens: t.List[Token], pos: int) -> bool:
+    try:
+        return (
+            tokens[pos].text.upper() == command.upper()
+            and tokens[pos + 1].token_type == TokenType.SEMICOLON
+        )
+    except IndexError:
+        return False
+
+
+JINJA_QUERY_BEGIN = "JINJA_QUERY_BEGIN"
+JINJA_STATEMENT_BEGIN = "JINJA_STATEMENT_BEGIN"
+JINJA_END = "JINJA_END"
+
+
+def _is_jinja_statement_begin(tokens: t.List[Token], pos: int) -> bool:
+    return _is_command_statement(JINJA_STATEMENT_BEGIN, tokens, pos)
+
+
+def _is_jinja_query_begin(tokens: t.List[Token], pos: int) -> bool:
+    return _is_command_statement(JINJA_QUERY_BEGIN, tokens, pos)
+
+
+def _is_jinja_end(tokens: t.List[Token], pos: int) -> bool:
+    return _is_command_statement(JINJA_END, tokens, pos)
+
+
+def jinja_query(query: str) -> JinjaQuery:
+    return JinjaQuery(this=exp.Literal.string(query))
+
+
+def jinja_statement(statement: str) -> JinjaStatement:
+    return JinjaStatement(this=exp.Literal.string(statement))
+
+
+class ChunkType(Enum):
+    JINJA_QUERY = auto()
+    JINJA_STATEMENT = auto()
+    SQL = auto()
+
+
 def parse(sql: str, default_dialect: t.Optional[str] = None) -> t.List[exp.Expression]:
     """Parse a sql string.
 
@@ -404,37 +455,51 @@ def parse(sql: str, default_dialect: t.Optional[str] = None) -> t.List[exp.Expre
     dialect = Dialect.get_or_raise(match.group(2) if match else default_dialect)()
 
     tokens = dialect.tokenizer.tokenize(sql)
-    chunks: t.List[t.Tuple[t.List, bool]] = [([], False)]
+    chunks: t.List[t.Tuple[t.List[Token], ChunkType]] = [([], ChunkType.SQL)]
     total = len(tokens)
 
-    for i, token in enumerate(tokens):
-        if token.token_type == TokenType.SEMICOLON:
-            if i < total - 1:
-                chunks.append(([], False))
+    pos = 0
+    while pos < total:
+        token = tokens[pos]
+        if _is_jinja_end(tokens, pos) or (
+            chunks[-1][1] == ChunkType.SQL
+            and token.token_type == TokenType.SEMICOLON
+            and pos < total - 1
+        ):
+            chunks.append(([], ChunkType.SQL))
+            if token.token_type == TokenType.SEMICOLON:
+                pos += 1
+            else:
+                # Jinja end statement
+                pos += 2
+        elif _is_jinja_query_begin(tokens, pos):
+            chunks.append(([], ChunkType.JINJA_QUERY))
+            pos += 2
+        elif _is_jinja_statement_begin(tokens, pos):
+            chunks.append(([], ChunkType.JINJA_STATEMENT))
+            pos += 2
         else:
-            if token.token_type == TokenType.BLOCK_START or (
-                i < total - 1
-                and token.token_type == TokenType.L_BRACE
-                and tokens[i + 1].token_type == TokenType.L_BRACE
-            ):
-                chunks[-1] = (chunks[-1][0], True)
             chunks[-1][0].append(token)
+            pos += 1
 
     expressions: t.List[exp.Expression] = []
 
-    for chunk, is_jinja in chunks:
-        if is_jinja:
+    for chunk, chunk_type in chunks:
+        if chunk_type == ChunkType.SQL:
+            for expression in dialect.parser().parse(chunk, sql):
+                if expression:
+                    expressions.append(expression)
+        else:
             start, *_, end = chunk
             segment = sql[start.start : end.end + 2]
             variables = [
                 exp.Literal.string(var)
                 for var in find_undeclared_variables(ENVIRONMENT.parse(segment))
             ]
-            expressions.append(Jinja(this=exp.Literal.string(segment), expressions=variables))
-        else:
-            for expression in dialect.parser().parse(chunk, sql):
-                if expression:
-                    expressions.append(expression)
+            klass = JinjaQuery if chunk_type == ChunkType.JINJA_QUERY else JinjaStatement
+            expressions.append(
+                klass(this=exp.Literal.string(segment.strip()), expressions=variables)
+            )
 
     return expressions
 
@@ -462,6 +527,8 @@ def extend_sqlglot() -> None:
                     MacroVar: lambda self, e: f"@{e.name}",
                     Model: _model_sql,
                     Jinja: lambda self, e: e.name,
+                    JinjaQuery: lambda self, e: f"{JINJA_QUERY_BEGIN};\n{e.name}\n{JINJA_END};",
+                    JinjaStatement: lambda self, e: f"{JINJA_STATEMENT_BEGIN};\n{e.name.strip()}\n{JINJA_END};",
                     ModelKind: _model_kind_sql,
                     PythonCode: lambda self, e: self.expressions(e, sep="\n", indent=False),
                 }

sqlmesh/core/model/definition.py

@@ -12,7 +12,7 @@
 
 from astor import to_source
 from pydantic import Field
-from sqlglot import diff, exp, parse_one
+from sqlglot import diff, exp
 from sqlglot.diff import Insert, Keep
 from sqlglot.helper import ensure_list
 from sqlglot.optimizer.scope import traverse_scope
@@ -26,7 +26,7 @@
 from sqlmesh.core.model.meta import ModelMeta
 from sqlmesh.core.model.seed import Seed, create_seed
 from sqlmesh.core.renderer import ExpressionRenderer, QueryRenderer
-from sqlmesh.utils.date import TimeLike, date_dict, make_inclusive, to_datetime
+from sqlmesh.utils.date import TimeLike, make_inclusive, to_datetime
 from sqlmesh.utils.errors import ConfigError, SQLMeshError, raise_config_error
 from sqlmesh.utils.jinja import JinjaMacroRegistry, extract_macro_references
 from sqlmesh.utils.metaprogramming import (
@@ -630,7 +630,7 @@ class SqlModel(_SqlBasedModel):
         post_statements: The list of SQL statements that follow after the model's query.
     """
 
-    query: t.Union[exp.Subqueryable, d.Jinja]
+    query: t.Union[exp.Subqueryable, d.JinjaQuery]
     source_type: Literal["sql"] = "sql"
 
     _columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None
@@ -1042,7 +1042,7 @@ def load_model(
 
     # Extract the query and any pre/post statements
     query_or_seed_insert, pre_statements, post_statements = _split_sql_model_statements(
-        expressions[1:], jinja_macros or JinjaMacroRegistry(), dialect, path
+        expressions[1:], path
     )
 
     meta_fields: t.Dict[str, t.Any] = {
@@ -1067,7 +1067,9 @@ def load_model(
         for r in references
     }
 
-    if query_or_seed_insert is not None and isinstance(query_or_seed_insert, exp.Subqueryable):
+    if query_or_seed_insert is not None and isinstance(
+        query_or_seed_insert, (exp.Subqueryable, d.JinjaQuery)
+    ):
         macro_references.update(extract_macro_references(query_or_seed_insert.sql(dialect=dialect)))
         return create_sql_model(
             name,
@@ -1103,7 +1105,7 @@ def load_model(
             )
         except Exception:
             raise_config_error(
-                "The model definition must either have a SELECT query or a valid Seed kind",
+                "The model definition must either have a SELECT query, a JINJA_QUERY block, or a valid Seed kind",
                 path,
             )
             raise
@@ -1142,9 +1144,9 @@ def create_sql_model(
         dialect: The default dialect if no model dialect is configured.
             The format must adhere to Python's strftime codes.
     """
-    if not isinstance(query, (exp.Subqueryable, d.Jinja)):
+    if not isinstance(query, (exp.Subqueryable, d.JinjaQuery)):
         raise_config_error(
-            "A query is required and must be a SELECT or UNION statement",
+            "A query is required and must be a SELECT statement, a UNION statement, or a JINJA_QUERY block",
             path,
         )
 
@@ -1337,7 +1339,7 @@ def _create_model(
 
 
 def _split_sql_model_statements(
-    expressions: t.List[exp.Expression], jinja_macros: JinjaMacroRegistry, dialect: str, path: Path
+    expressions: t.List[exp.Expression], path: Path
 ) -> t.Tuple[t.Optional[exp.Expression], t.List[exp.Expression], t.List[exp.Expression]]:
     """Extracts the SELECT query from a sequence of expressions.
 
@@ -1353,15 +1355,10 @@ def _split_sql_model_statements(
     """
     query_positions = []
     for idx, expression in enumerate(expressions):
-        # Render the expression using the macro registry if the expression is jinja.
-        if isinstance(expression, d.Jinja):
-            jinja_env = jinja_macros.build_environment(**date_dict(c.EPOCH, c.EPOCH, c.EPOCH))
-            rendered_expression = jinja_env.from_string(expression.name).render()
-            if rendered_expression is None:
-                continue
-            expression = parse_one(rendered_expression, read=dialect)
-
-        if isinstance(expression, exp.Subqueryable) or expression == INSERT_SEED_MACRO_CALL:
+        if (
+            isinstance(expression, (exp.Subqueryable, d.JinjaQuery))
+            or expression == INSERT_SEED_MACRO_CALL
+        ):
             query_positions.append((expression, idx))
 
     if not query_positions: