Repository: TobikoData/sqlmesh
Pull request #2499 (2024-04-23 23:06:37)

Feat: Support macros in the MODEL statement
This update allow users to easily parameterize their model names and other attributes by using macros and variables inside the `MODEL` statement. Thus the following example is now a valid definition:
```
MODEL (
    name @{gateway}.test_model
);
```


Patch:
diff --git a/.circleci/continue_config.yml b/.circleci/continue_config.yml
index a8a416af..137dd3ff 100644
--- a/.circleci/continue_config.yml
+++ b/.circleci/continue_config.yml
@@ -231,7 +231,8 @@ jobs:
         - run:
             name: Trigger private tests
             command: |
-              echo "export COMMIT_MESSAGE=\"$(git log --format=%s -n 1 $CIRCLE_SHA1)\"" >> "$BASH_ENV"
+              echo 'export COMMIT_MESSAGE="$(git log --format=%s -n 1 $CIRCLE_SHA1)"' >> "$BASH_ENV"
+              echo 'export FORMATTED_COMMIT_MESSAGE="${COMMIT_MESSAGE//\"/\\\"}"' >> "$BASH_ENV"
               source "$BASH_ENV"
               curl --request POST \
                 --url $TOBIKO_PRIVATE_CIRCLECI_URL \
@@ -245,7 +246,7 @@ jobs:
                     "sqlmesh_branch":"'$CIRCLE_BRANCH'",
                     "sqlmesh_commit_author":"'$CIRCLE_USERNAME'",
                     "sqlmesh_commit_hash":"'$CIRCLE_SHA1'",
-                    "sqlmesh_commit_message":"'"$COMMIT_MESSAGE"'"
+                    "sqlmesh_commit_message":"'"$FORMATTED_COMMIT_MESSAGE"'"
                     }
                 }'
 
diff --git a/sqlmesh/core/macros.py b/sqlmesh/core/macros.py
index c0ba8da4..0b52569a 100644
--- a/sqlmesh/core/macros.py
+++ b/sqlmesh/core/macros.py
@@ -235,7 +235,9 @@ class MacroEvaluator:
         """
         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
diff --git a/sqlmesh/core/model/definition.py b/sqlmesh/core/model/definition.py
index f247e84d..ab2d2801 100644
--- a/sqlmesh/core/model/definition.py
+++ b/sqlmesh/core/model/definition.py
@@ -22,7 +22,7 @@ from sqlglot.time import format_time
 
 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:
diff --git a/sqlmesh/core/state_sync/engine_adapter.py b/sqlmesh/core/state_sync/engine_adapter.py
index b2960aed..82bc6a9f 100644
--- a/sqlmesh/core/state_sync/engine_adapter.py
+++ b/sqlmesh/core/state_sync/engine_adapter.py
@@ -39,6 +39,7 @@ from sqlmesh.core.snapshot import (
     Intervals,
     Node,
     Snapshot,
+    SnapshotChangeCategory,
     SnapshotFingerprint,
     SnapshotId,
     SnapshotIdLike,
@@ -215,7 +216,10 @@ class EngineAdapterStateSync(CommonStateSyncMixin, StateSync):
         snapshots_to_store = []
 
         for snapshot in snapshots:
-            if isinstance(snapshot.node, SeedModel):
+            if isinstance(snapshot.node, SeedModel) and snapshot.change_category in (
+                SnapshotChangeCategory.BREAKING,
+                SnapshotChangeCategory.NON_BREAKING,
+            ):
                 seed_model = t.cast(SeedModel, snapshot.node)
                 seed_contents.append(
                     {


Test patch:
diff --git a/tests/core/test_model.py b/tests/core/test_model.py
index 88e524ae..4faf266e 100644
--- a/tests/core/test_model.py
+++ b/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"
diff --git a/tests/core/test_state_sync.py b/tests/core/test_state_sync.py
index c0d9c032..f4fef1fe 100644
--- a/tests/core/test_state_sync.py
+++ b/tests/core/test_state_sync.py
@@ -2225,3 +2225,40 @@ def test_snapshot_batching(state_sync, mocker, make_snapshot):
     assert len(snapshots) == 3
     calls = mock.fetchall.call_args_list
     assert len(calls) == 2
+
+
+def test_seed_model_metadata_update(
+    state_sync: EngineAdapterStateSync,
+    make_snapshot: t.Callable,
+):
+    model = SeedModel(
+        name="a",
+        kind=SeedKind(path="./path/to/seed"),
+        seed=Seed(content="header\n1\n2"),
+        column_hashes={"header": "hash"},
+        depends_on=set(),
+    )
+    snapshot = make_snapshot(model)
+    snapshot.categorize_as(SnapshotChangeCategory.BREAKING)
+
+    state_sync.push_snapshots([snapshot])
+
+    seed_query = (
+        exp.select("COUNT(*)")
+        .from_("sqlmesh._seeds")
+        .where(exp.column("version").eq(snapshot.version))
+    )
+
+    assert state_sync.engine_adapter.fetchone(seed_query)[0] == 1
+
+    model = model.copy(update={"owner": "jen"})
+    new_snapshot = make_snapshot(model)
+    new_snapshot.previous_versions = snapshot.all_versions
+    new_snapshot.categorize_as(SnapshotChangeCategory.FORWARD_ONLY)
+
+    assert snapshot.fingerprint != new_snapshot.fingerprint
+    assert snapshot.version == new_snapshot.version
+
+    state_sync.push_snapshots([new_snapshot])
+    assert state_sync.engine_adapter.fetchone(seed_query)[0] == 1
+    assert len(state_sync.get_snapshots([new_snapshot, snapshot])) == 2