Feat: Allow macros in python model properties (#3740)

Author: themisvaltinos

sqlmesh/core/model/decorator.py

@@ -17,6 +17,7 @@
     create_python_model,
     create_sql_model,
     get_model_name,
+    render_meta_fields,
 )
 from sqlmesh.core.model.kind import ModelKindName, _ModelKind
 from sqlmesh.utils import registry_decorator
@@ -118,6 +119,21 @@ def model(
 
         build_env(self.func, env=env, name=entrypoint, path=module_path)
 
+        rendered_fields = render_meta_fields(
+            fields={"name": self.name, **self.kwargs},
+            module_path=module_path,
+            macros=macros,
+            jinja_macros=jinja_macros,
+            variables=variables,
+            path=path,
+            dialect=dialect,
+            default_catalog=default_catalog,
+        )
+
+        rendered_name = rendered_fields["name"]
+        if isinstance(rendered_name, exp.Expression):
+            rendered_fields["name"] = rendered_name.sql(dialect=dialect)
+
         common_kwargs = {
             "defaults": defaults,
             "path": path,
@@ -133,7 +149,7 @@ def model(
             "macros": macros,
             "jinja_macros": jinja_macros,
             "audit_definitions": audit_definitions,
-            **self.kwargs,
+            **rendered_fields,
         }
 
         for key in ("pre_statements", "post_statements", "on_virtual_update"):
@@ -146,5 +162,5 @@ def model(
 
         if self.is_sql:
             query = MacroFunc(this=exp.Anonymous(this=entrypoint))
-            return create_sql_model(self.name, query, **common_kwargs)
-        return create_python_model(self.name, entrypoint, **common_kwargs)
+            return create_sql_model(query=query, **common_kwargs)
+        return create_python_model(entrypoint=entrypoint, **common_kwargs)

sqlmesh/core/model/definition.py

@@ -62,6 +62,15 @@
 
 logger = logging.getLogger(__name__)
 
+RUNTIME_RENDERED_MODEL_FIELDS = {
+    "audits",
+    "signals",
+    "description",
+    "cron",
+    "physical_properties",
+    "merge_filter",
+}
+
 
 class _Model(ModelMeta, frozen=True):
     """Model is the core abstraction for user defined datasets.
@@ -1823,7 +1832,7 @@ def load_sql_based_model(
                 if kind_prop.name.lower() == "merge_filter":
                     unrendered_merge_filter = kind_prop
 
-    meta_renderer = _meta_renderer(
+    rendered_meta_exprs = render_expression(
         expression=meta,
         module_path=module_path,
         macros=macros,
@@ -1834,7 +1843,6 @@ def load_sql_based_model(
         default_catalog=default_catalog,
     )
 
-    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)}",
@@ -2024,21 +2032,6 @@ def create_python_model(
     # Also remove self-references that are found
 
     dialect = kwargs.get("dialect")
-    renderer_kwargs = {
-        "module_path": module_path,
-        "macros": macros,
-        "jinja_macros": jinja_macros,
-        "variables": variables,
-        "path": path,
-        "dialect": dialect,
-        "default_catalog": kwargs.get("default_catalog"),
-    }
-
-    name_renderer = _meta_renderer(
-        expression=d.parse_one(name, dialect=dialect),
-        **renderer_kwargs,  # type: ignore
-    )
-    name = t.cast(t.List[exp.Expression], name_renderer.render())[0].sql(dialect=dialect)
 
     dependencies_unspecified = depends_on is None
 
@@ -2050,15 +2043,21 @@ def create_python_model(
     if dependencies_unspecified:
         depends_on = parsed_depends_on - {name}
     else:
-        depends_on_renderer = _meta_renderer(
+        depends_on_rendered = render_expression(
             expression=exp.Array(
                 expressions=[d.parse_one(dep, dialect=dialect) for dep in depends_on or []]
             ),
-            **renderer_kwargs,  # type: ignore
+            module_path=module_path,
+            macros=macros,
+            jinja_macros=jinja_macros,
+            variables=variables,
+            path=path,
+            dialect=dialect,
+            default_catalog=kwargs.get("default_catalog"),
         )
         depends_on = {
             dep.sql(dialect=dialect)
-            for dep in t.cast(t.List[exp.Expression], depends_on_renderer.render())[0].expressions
+            for dep in t.cast(t.List[exp.Expression], depends_on_rendered)[0].expressions
         }
 
     variables = {k: v for k, v in (variables or {}).items() if k in referenced_variables}
@@ -2382,7 +2381,61 @@ def _refs_to_sql(values: t.Any) -> exp.Expression:
     return exp.Tuple(expressions=values)
 
 
-def _meta_renderer(
+def render_meta_fields(
+    fields: t.Dict[str, t.Any],
+    module_path: Path,
+    path: Path,
+    jinja_macros: t.Optional[JinjaMacroRegistry],
+    macros: t.Optional[MacroRegistry],
+    dialect: DialectType,
+    variables: t.Optional[t.Dict[str, t.Any]],
+    default_catalog: t.Optional[str],
+) -> t.Dict[str, t.Any]:
+    def render_field_value(value: t.Any) -> t.Any:
+        if isinstance(value, exp.Expression) or (
+            isinstance(value, str) and d.SQLMESH_MACRO_PREFIX in value
+        ):
+            expression = exp.maybe_parse(value, dialect=dialect)
+            rendered_expr = render_expression(
+                expression=expression,
+                module_path=module_path,
+                macros=macros,
+                jinja_macros=jinja_macros,
+                variables=variables,
+                path=path,
+                dialect=dialect,
+                default_catalog=default_catalog,
+            )
+            if rendered_expr is None:
+                raise SQLMeshError(
+                    f"Failed to render model attribute `{fields['name']}` at `{path}`\n"
+                    f"'{expression.sql(dialect=dialect)}' must return an expression"
+                )
+            if len(rendered_expr) != 1:
+                raise SQLMeshError(
+                    f"Failed to render model attribute `{fields['name']}` at `{path}`.\n"
+                    f"`{expression.sql(dialect=dialect)}` must return one result, but got {len(rendered_expr)}"
+                )
+            return rendered_expr[0]
+
+        return value
+
+    for field_name, field_info in ModelMeta.all_field_infos().items():
+        field = field_info.alias or field_name
+        if field not in RUNTIME_RENDERED_MODEL_FIELDS and (field_value := fields.get(field)):
+            if isinstance(field_value, dict):
+                for key in list(field_value.keys()):
+                    if key not in RUNTIME_RENDERED_MODEL_FIELDS:
+                        fields[field][key] = render_field_value(field_value[key])
+            elif isinstance(field_value, list):
+                fields[field] = [render_field_value(value) for value in field_value]
+            else:
+                fields[field] = render_field_value(field_value)
+
+    return fields
+
+
+def render_expression(
     expression: exp.Expression,
     module_path: Path,
     path: Path,
@@ -2391,7 +2444,7 @@ def _meta_renderer(
     dialect: DialectType = None,
     variables: t.Optional[t.Dict[str, t.Any]] = None,
     default_catalog: t.Optional[str] = None,
-) -> ExpressionRenderer:
+) -> t.Optional[t.List[exp.Expression]]:
     meta_python_env = make_python_env(
         expressions=expression,
         jinja_macro_references=None,
@@ -2410,7 +2463,7 @@ def _meta_renderer(
         default_catalog=default_catalog,
         quote_identifiers=False,
         normalize_identifiers=False,
-    )
+    ).render()
 
 
 META_FIELD_CONVERTER: t.Dict[str, t.Callable] = {

tests/core/test_context.py

@@ -1054,6 +1054,27 @@ def test_disabled_model(copy_to_temp_path):
     assert not context.get_model("sushi.disabled_py")
 
 
+def test_disabled_model_python_macro(sushi_context):
+    @model(
+        "memory.sushi.disabled_model_2",
+        columns={"col": "int"},
+        enabled="@IF(@gateway = 'dev', True, False)",
+    )
+    def entrypoint(context, **kwargs):
+        yield pd.DataFrame({"col": []})
+
+    test_model = model.get_registry()["memory.sushi.disabled_model_2"].model(
+        module_path=Path("."), path=Path("."), variables={"gateway": "prod"}
+    )
+    assert not test_model.enabled
+
+    with pytest.raises(
+        SQLMeshError,
+        match="The disabled model 'memory.sushi.disabled_model_2' cannot be upserted",
+    ):
+        sushi_context.upsert_model(test_model)
+
+
 def test_get_model_mixed_dialects(copy_to_temp_path):
     path = copy_to_temp_path("examples/sushi")