Repository: TobikoData/sqlmesh
Pull request #2833 (2024-06-27 00:45:02)

Feat: Introduce metadata only flag for macro functions
This PR allows flagging macro functions in a SQL model's pre and post hooks to be included in the model's metadata hash instead of its data hash.

tobymao: why does it need to be a flag? why isn’t it always metadata?
vchan: > why does it need to be a flag? why isn’t it always metadata?

@tobymao Good question. I had a chat with @izeigerman and we felt it'd be safer to use a flag. A macro func could do something to warrant a backfill even though it doesn't render any expressions to execute. Happy to discuss this more though. I'm not that familiar with the macros system.
tobymao: > > why does it need to be a flag? why isn’t it always metadata?
> 
> @tobymao Good question. I had a chat with @izeigerman and we felt it'd be safer to use a flag. A macro func could do something to warrant a backfill even though it doesn't render any expressions to execute. Happy to discuss this more though. I'm not that familiar with the macros system.

ok that's fine with me, but we should edit the docs to indicate this exists
vchan: > > > why does it need to be a flag? why isn’t it always metadata?
> > 
> > 
> > @tobymao Good question. I had a chat with @izeigerman and we felt it'd be safer to use a flag. A macro func could do something to warrant a backfill even though it doesn't render any expressions to execute. Happy to discuss this more though. I'm not that familiar with the macros system.
> 
> ok that's fine with me, but we should edit the docs to indicate this exists
> 
> so the default is a data hash or metadata hash? default should be metadata right?

Okay, I'll update the docs. As for the default, it's currently data hash since it's safer (will trigger a backfill for your plan vs wait until next run). What do you think?

Patch:
diff --git a/sqlmesh/core/constants.py b/sqlmesh/core/constants.py
index 4cc897b4..593075ac 100644
--- a/sqlmesh/core/constants.py
+++ b/sqlmesh/core/constants.py
@@ -60,6 +60,7 @@ GATEWAY = "gateway"
 
 SQLMESH_MACRO = "__sqlmesh__macro__"
 SQLMESH_BUILTIN = "__sqlmesh__builtin__"
+SQLMESH_METADATA = "__sqlmesh__metadata__"
 
 
 BUILTIN = "builtin"
diff --git a/sqlmesh/core/macros.py b/sqlmesh/core/macros.py
index f96f4262..629e03b1 100644
--- a/sqlmesh/core/macros.py
+++ b/sqlmesh/core/macros.py
@@ -558,9 +558,15 @@ class macro(registry_decorator):
 
     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.
diff --git a/sqlmesh/core/model/definition.py b/sqlmesh/core/model/definition.py
index d08fb312..e4fb9345 100644
--- a/sqlmesh/core/model/definition.py
+++ b/sqlmesh/core/model/definition.py
@@ -962,11 +962,11 @@ class _SqlBasedModel(_Model):
 
         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]
@@ -974,6 +974,28 @@ class _SqlBasedModel(_Model):
 
         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.
diff --git a/sqlmesh/utils/metaprogramming.py b/sqlmesh/utils/metaprogramming.py
index ac25bcc7..786e847d 100644
--- a/sqlmesh/utils/metaprogramming.py
+++ b/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(


Test patch:
diff --git a/tests/core/test_macros.py b/tests/core/test_macros.py
index 2262fd87..a1786956 100644
--- a/tests/core/test_macros.py
+++ b/tests/core/test_macros.py
@@ -3,6 +3,7 @@ import typing as t
 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
diff --git a/tests/core/test_model.py b/tests/core/test_model.py
index 29433495..1c45592c 100644
--- a/tests/core/test_model.py
+++ b/tests/core/test_model.py
@@ -1616,7 +1616,7 @@ FROM (VALUES
 WHERE
   FALSE
 LIMIT 0
-	""",
+""",
     )
 
 
@@ -5096,3 +5096,50 @@ materialized FALSE
 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={})
diff --git a/tests/schedulers/airflow/test_client.py b/tests/schedulers/airflow/test_client.py
index e30723dd..14e025ea 100644
--- a/tests/schedulers/airflow/test_client.py
+++ b/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")]},
diff --git a/tests/utils/test_metaprogramming.py b/tests/utils/test_metaprogramming.py
index 87aeda8e..f236a752 100644
--- a/tests/utils/test_metaprogramming.py
+++ b/tests/utils/test_metaprogramming.py
@@ -9,6 +9,7 @@ from pytest_mock.plugin import MockerFixture
 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 @@ class DataClass:
             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",