Feat: Introduce metadata only flag for macro functions (#2833)

Author: vchan

sqlmesh/core/constants.py

@@ -60,6 +60,7 @@
 
 SQLMESH_MACRO = "__sqlmesh__macro__"
 SQLMESH_BUILTIN = "__sqlmesh__builtin__"
+SQLMESH_METADATA = "__sqlmesh__metadata__"
 
 
 BUILTIN = "builtin"

sqlmesh/core/macros.py

@@ -558,9 +558,15 @@ def add_one(evaluator: MacroEvaluator, column: exp.Literal) -> exp.Add:
 
     registry_name = "macros"
 
+    def __init__(self, *args: t.Any, is_metadata: bool = False, **kwargs: t.Any) -> None:
+        super().__init__(*args, **kwargs)
+        self.is_metadata = is_metadata
+
     def __call__(
         self, func: t.Callable[..., DECORATOR_RETURN_TYPE]
     ) -> t.Callable[..., DECORATOR_RETURN_TYPE]:
+        if self.is_metadata:
+            setattr(func, c.SQLMESH_METADATA, self.is_metadata)
         wrapper = super().__call__(func)
 
         # This is used to identify macros at runtime to unwrap during serialization.

sqlmesh/core/model/definition.py

@@ -962,18 +962,40 @@ def _data_hash_values(self) -> t.List[str]:
 
         for statement in (*self.pre_statements, *self.post_statements):
             statement_exprs: t.List[exp.Expression] = []
-            if isinstance(statement, d.MacroDef):
-                statement_exprs = [statement]
-            else:
+            if not isinstance(statement, d.MacroDef):
                 rendered = self._statement_renderer(statement).render()
-                if rendered is not None:
+                if self._is_metadata_statement(statement):
+                    continue
+                if rendered:
                     statement_exprs = rendered
                 else:
                     statement_exprs = [statement]
             data_hash_values.extend(gen(e) for e in statement_exprs)
 
         return data_hash_values
 
+    @property
+    def _additional_metadata(self) -> t.List[str]:
+        additional_metadata = super()._additional_metadata
+
+        for statement in (*self.pre_statements, *self.post_statements):
+            if self._is_metadata_statement(statement):
+                additional_metadata.append(gen(statement))
+
+        return additional_metadata
+
+    def _is_metadata_statement(self, statement: exp.Expression) -> bool:
+        if isinstance(statement, d.MacroDef):
+            return True
+        if isinstance(statement, d.MacroFunc):
+            target_macro = macro.get_registry().get(statement.name)
+            if target_macro:
+                return target_macro.is_metadata
+            target_macro = self.python_env.get(statement.name)
+            if isinstance(target_macro, Executable):
+                return bool(target_macro.is_metadata)
+        return False
+
 
 class SqlModel(_SqlBasedModel):
     """The model definition which relies on a SQL query to fetch the data.

sqlmesh/utils/metaprogramming.py

@@ -314,6 +314,7 @@ class Executable(PydanticModel):
     name: t.Optional[str] = None
     path: t.Optional[str] = None
     alias: t.Optional[str] = None
+    is_metadata: t.Optional[bool] = None
 
     @property
     def is_definition(self) -> bool:
@@ -360,6 +361,7 @@ def serialize_env(env: t.Dict[str, t.Any], path: Path) -> t.Dict[str, Executable
                     # Do `as_posix` to serialize windows path back to POSIX
                     path=t.cast(Path, file_path).relative_to(path.absolute()).as_posix(),
                     alias=k if name != k else None,
+                    is_metadata=getattr(v, c.SQLMESH_METADATA, None),
                 )
             else:
                 serialized[k] = Executable(

tests/core/test_macros.py

@@ -3,6 +3,7 @@
 import pytest
 from sqlglot import MappingSchema, ParseError, exp, parse_one
 
+from sqlmesh.core import constants as c
 from sqlmesh.core.dialect import StagedFilePath
 from sqlmesh.core.macros import SQL, MacroEvalError, MacroEvaluator, macro
 from sqlmesh.utils.errors import SQLMeshError
@@ -599,3 +600,19 @@ def test_macro_parameter_resolution(macro_evaluator):
     with pytest.raises(MacroEvalError) as e:
         macro_evaluator.evaluate(parse_one("@test_arg_resolution(1, 2, 3)"))
     assert str(e.value.__cause__) == "too many positional arguments"
+
+
+def test_macro_metadata_flag():
+    @macro()
+    def noop(evaluator) -> None:
+        return None
+
+    @macro(is_metadata=True)
+    def noop_metadata_only(evaluator) -> None:
+        return None
+
+    assert not hasattr(noop, c.SQLMESH_METADATA)
+    assert not macro.get_registry()["noop"].is_metadata
+
+    assert getattr(noop_metadata_only, c.SQLMESH_METADATA) is True
+    assert macro.get_registry()["noop_metadata_only"].is_metadata is True

tests/core/test_model.py

@@ -1616,7 +1616,7 @@ def my_model(context, **kwargs):
 WHERE
   FALSE
 LIMIT 0
-	""",
+""",
     )
 
 
@@ -5096,3 +5096,50 @@ def test_model_kind_to_expression():
 materialized TRUE
 )"""
     )
+
+
+@pytest.mark.parametrize(
+    "is_metadata",
+    [True, False],
+)
+def test_macro_func_hash(is_metadata):
+    macro.set_registry({})
+
+    @macro(is_metadata=is_metadata)
+    def noop(evaluator) -> None:
+        return None
+
+    # Prevent macro from entering python_env so data hash is consistent
+    setattr(macro.get_registry()["noop"], c.SQLMESH_BUILTIN, True)
+
+    expressions = d.parse(
+        """
+        MODEL (
+            name db.model,
+        );
+
+        SELECT 1;
+    """
+    )
+    model = load_sql_based_model(expressions, path=Path("./examples/sushi/models/test_model.sql"))
+
+    expressions = d.parse(
+        """
+        MODEL (
+            name db.model,
+        );
+
+        SELECT 1;
+
+        @noop();
+    """
+    )
+    new_model = load_sql_based_model(
+        expressions, path=Path("./examples/sushi/models/test_model.sql")
+    )
+    if is_metadata:
+        assert model.data_hash == new_model.data_hash
+        assert model.metadata_hash(audits={}) != new_model.metadata_hash(audits={})
+    else:
+        assert model.data_hash != new_model.data_hash
+        assert model.metadata_hash(audits={}) == new_model.metadata_hash(audits={})

tests/schedulers/airflow/test_client.py

@@ -174,7 +174,7 @@ def test_apply_plan(mocker: MockerFixture, snapshot: Snapshot):
         "models_to_backfill": ['"test_model"'],
         "end_bounded": False,
         "ensure_finalized_snapshots": False,
-        "directly_modified_snapshots": [{"identifier": "349602315", "name": '"test_model"'}],
+        "directly_modified_snapshots": [{"identifier": "844700562", "name": '"test_model"'}],
         "indirectly_modified_snapshots": {},
         "removed_snapshots": [],
         "restatements": {'"test_model"': [to_timestamp("2024-01-01"), to_timestamp("2024-01-02")]},

tests/utils/test_metaprogramming.py

@@ -9,6 +9,7 @@
 from sqlglot.expressions import to_table
 
 import tests.utils.test_date as test_date
+from sqlmesh.core import constants as c
 from sqlmesh.utils.errors import SQLMeshError
 from sqlmesh.utils.metaprogramming import (
     Executable,
@@ -41,7 +42,7 @@ def test_print_exception(mocker: MockerFixture):
 
     expected_message = f"""Traceback (most recent call last):
 
-  File "{__file__}", line 38, in test_print_exception
+  File "{__file__}", line 39, in test_print_exception
     eval("test_fun()", env)
 
   File "<string>", line 1, in <module>
@@ -95,11 +96,19 @@ def other_func(a: int) -> int:
     return X + a
 
 
+def noop_metadata() -> None:
+    return None
+
+
+setattr(noop_metadata, c.SQLMESH_METADATA, True)
+
+
 def main_func(y: int) -> int:
     """DOC STRING"""
     sqlglot.parse_one("1")
     MyClass()
     DataClass(x=y)
+    noop_metadata()
 
     def closure(z: int) -> int:
         return z + Z
@@ -113,6 +122,7 @@ def test_func_globals() -> None:
         "Z": 3,
         "DataClass": DataClass,
         "MyClass": MyClass,
+        "noop_metadata": noop_metadata,
         "other_func": other_func,
         "sqlglot": sqlglot,
     }
@@ -144,6 +154,7 @@ def test_normalize_source() -> None:
     sqlglot.parse_one('1')
     MyClass()
     DataClass(x=y)
+    noop_metadata()
 
     def closure(z: int):
         return z + Z
@@ -183,6 +194,7 @@ def test_serialize_env() -> None:
     sqlglot.parse_one('1')
     MyClass()
     DataClass(x=y)
+    noop_metadata()
 
     def closure(z: int):
         return z + Z
@@ -233,6 +245,13 @@ def baz(self):
             path="test_metaprogramming.py",
             payload="my_lambda = lambda : print('z')",
         ),
+        "noop_metadata": Executable(
+            name="noop_metadata",
+            path="test_metaprogramming.py",
+            payload="""def noop_metadata():
+    return None""",
+            is_metadata=True,
+        ),
         "other_func": Executable(
             name="other_func",
             path="test_metaprogramming.py",