Feat: Support macros in the MODEL statement (#2499)

Author: izeigerman

sqlmesh/core/macros.py

@@ -235,7 +235,9 @@ def template(self, text: t.Any, local_variables: t.Dict[str, t.Any]) -> str:
         """
         mapping = {}
 
-        for k, v in chain(self.locals.items(), local_variables.items()):
+        variables = self.locals.get(c.SQLMESH_VARS, {})
+
+        for k, v in chain(variables.items(), self.locals.items(), local_variables.items()):
             # try to convert all variables into sqlglot expressions
             # because they're going to be converted into strings in sql
             # we use bare Exception instead of ValueError because there's

sqlmesh/core/model/definition.py

@@ -22,7 +22,7 @@
 
 from sqlmesh.core import constants as c
 from sqlmesh.core import dialect as d
-from sqlmesh.core.macros import MacroRegistry, macro
+from sqlmesh.core.macros import MacroRegistry, MacroStrTemplate, macro
 from sqlmesh.core.model.common import expression_validator
 from sqlmesh.core.model.kind import ModelKindName, SeedKind
 from sqlmesh.core.model.meta import ModelMeta
@@ -1435,6 +1435,38 @@ def load_sql_based_model(
             path,
         )
 
+    unrendered_signals = None
+    for prop in meta.expressions:
+        if prop.name.lower() == "signals":
+            unrendered_signals = prop.args.get("value")
+
+    meta_python_env = _python_env(
+        expressions=meta,
+        jinja_macro_references=None,
+        module_path=module_path,
+        macros=macros or macro.get_registry(),
+        variables=variables,
+        path=path,
+    )
+    meta_renderer = ExpressionRenderer(
+        meta,
+        dialect,
+        [],
+        path=path,
+        jinja_macro_registry=jinja_macros,
+        python_env=meta_python_env,
+        default_catalog=default_catalog,
+        quote_identifiers=False,
+    )
+    rendered_meta_exprs = meta_renderer.render()
+    if rendered_meta_exprs is None or len(rendered_meta_exprs) != 1:
+        raise_config_error(
+            f"Invalid MODEL statement:\n{meta.sql(dialect=dialect, pretty=True)}",
+            path,
+        )
+        raise
+    rendered_meta = rendered_meta_exprs[0]
+
     # Extract the query and any pre/post statements
     query_or_seed_insert, pre_statements, post_statements = _split_sql_model_statements(
         expressions[1:], path
@@ -1443,11 +1475,16 @@ def load_sql_based_model(
     meta_fields: t.Dict[str, t.Any] = {
         "dialect": dialect,
         "description": (
-            "\n".join(comment.strip() for comment in meta.comments) if meta.comments else None
+            "\n".join(comment.strip() for comment in rendered_meta.comments)
+            if rendered_meta.comments
+            else None
         ),
-        **{prop.name.lower(): prop.args.get("value") for prop in meta.expressions},
+        **{prop.name.lower(): prop.args.get("value") for prop in rendered_meta.expressions},
         **kwargs,
     }
+    if unrendered_signals:
+        # Signals must remain unrendered, so that they can be rendered later at evaluation runtime.
+        meta_fields["signals"] = unrendered_signals
 
     name = meta_fields.pop("name", "")
     if not name:
@@ -1465,8 +1502,8 @@ def load_sql_based_model(
     )
 
     jinja_macros = (jinja_macros or JinjaMacroRegistry()).trim(jinja_macro_references)
-    for macro in jinja_macros.root_macros.values():
-        used_variables.update(extract_macro_references_and_variables(macro.definition)[1])
+    for jinja_macro in jinja_macros.root_macros.values():
+        used_variables.update(extract_macro_references_and_variables(jinja_macro.definition)[1])
 
     common_kwargs = dict(
         pre_statements=pre_statements,
@@ -1867,7 +1904,7 @@ def _python_env(
     expressions = ensure_list(expressions)
     for expression in expressions:
         if not isinstance(expression, d.Jinja):
-            for macro_func_or_var in expression.find_all(d.MacroFunc, d.MacroVar):
+            for macro_func_or_var in expression.find_all(d.MacroFunc, d.MacroVar, exp.Identifier):
                 if macro_func_or_var.__class__ is d.MacroFunc:
                     name = macro_func_or_var.this.name.lower()
                     if name in macros:
@@ -1888,6 +1925,15 @@ def _python_env(
                         used_macros[name] = macros[name]
                     elif name in variables:
                         used_variables.add(name)
+                elif (
+                    isinstance(macro_func_or_var, exp.Identifier) and "@" in macro_func_or_var.this
+                ):
+                    for _, identifier, braced_identifier, _ in MacroStrTemplate.pattern.findall(
+                        macro_func_or_var.this
+                    ):
+                        var_name = braced_identifier or identifier
+                        if var_name in variables:
+                            used_variables.add(var_name)
 
     for macro_ref in jinja_macro_references or set():
         if macro_ref.package is None and macro_ref.name in macros:

tests/core/test_model.py

@@ -938,7 +938,7 @@ def test_render_definition():
             owner owner_name,
             dialect spark,
             kind INCREMENTAL_BY_TIME_RANGE (
-                time_column (a, 'yyyymmdd')
+                time_column (`a`, 'yyyymmdd')
             ),
             storage_format iceberg,
             partitioned_by `a`,
@@ -3516,6 +3516,56 @@ def test_named_variable_macros() -> None:
     )
 
 
+def test_variables_in_templates() -> None:
+    model = load_sql_based_model(
+        parse(
+            """
+        MODEL(name sushi.test_gateway_macro);
+        @DEF(overridden_var, overridden_value);
+        SELECT 'gateway' AS col_@gateway, 'test_var_a' AS @{test_var_a}_col, 'overridden_var' AS col_@{overridden_var}_col
+        """
+        ),
+        variables={
+            c.GATEWAY: "in_memory",
+            "test_var_a": "test_value",
+            "test_var_unused": "unused",
+            "overridden_var": "initial_value",
+        },
+    )
+
+    assert model.python_env[c.SQLMESH_VARS] == Executable.value(
+        {c.GATEWAY: "in_memory", "test_var_a": "test_value", "overridden_var": "initial_value"}
+    )
+    assert (
+        model.render_query_or_raise().sql()
+        == "SELECT 'gateway' AS \"col_in_memory\", 'test_var_a' AS \"test_value_col\", 'overridden_var' AS \"col_overridden_value_col\""
+    )
+
+    model = load_sql_based_model(
+        parse(
+            """
+        MODEL(name sushi.test_gateway_macro);
+        @DEF(overridden_var, overridden_value);
+        SELECT 'combo' AS col_@{test_var_a}_@{overridden_var}_col_@gateway
+        """
+        ),
+        variables={
+            c.GATEWAY: "in_memory",
+            "test_var_a": "test_value",
+            "test_var_unused": "unused",
+            "overridden_var": "initial_value",
+        },
+    )
+
+    assert model.python_env[c.SQLMESH_VARS] == Executable.value(
+        {c.GATEWAY: "in_memory", "test_var_a": "test_value", "overridden_var": "initial_value"}
+    )
+    assert (
+        model.render_query_or_raise().sql()
+        == "SELECT 'combo' AS \"col_test_value_overridden_value_col_in_memory\""
+    )
+
+
 def test_variables_jinja():
     expressions = parse(
         """
@@ -3759,3 +3809,28 @@ def test_this_model() -> None:
 
     assert model.render_pre_statements()[0].sql() == """COPY "db"."table" TO 'a'"""
     assert model.render_post_statements()[0].sql() == """COPY "db"."table" TO 'b'"""
+
+
+def test_macros_in_model_statement(sushi_context, assert_exp_eq):
+    expressions = d.parse(
+        """
+        MODEL (
+            name @{gateway}__@{gateway}.test_model,
+            kind INCREMENTAL_BY_TIME_RANGE (
+                time_column @{time_column}
+
+            ),
+            start @IF(@gateway = 'test_gateway', '2023-01-01', '2024-01-02')
+        );
+
+        SELECT a, b UNION SELECT c, c
+        """
+    )
+
+    model = load_sql_based_model(
+        expressions, variables={"gateway": "test_gateway", "time_column": "a"}
+    )
+    assert model.name == "test_gateway__test_gateway.test_model"
+    assert model.time_column
+    assert model.time_column.column == exp.column("a", quoted=True)
+    assert model.start == "2023-01-01"