Fix!: make metadata_only a transitive property in Python objects

Author: georgesittas

sqlmesh/utils/metaprogramming.py

@@ -283,13 +283,25 @@ def build_env(
     # We don't rely on `env` to keep track of visited objects, because it's populated in post-order
     visited: t.Set[str] = set()
 
-    def walk(obj: t.Any, name: str) -> None:
+    def walk(obj: t.Any, name: str, is_metadata: t.Optional[bool] = None) -> None:
         obj_module = inspect.getmodule(obj)
         if name in visited or (obj_module and obj_module.__name__ == "builtins"):
             return
 
         visited.add(name)
-        if name not in env:
+        name_missing_from_env = name not in env
+
+        if name_missing_from_env or (
+            not is_metadata and env[name] == obj and getattr(env[name], c.SQLMESH_METADATA, None)
+        ):
+            if not name_missing_from_env:
+                # The existing object in the env is "metadata only" but we're walking it again as a
+                # non-"metadata only" dependency, so we update this flag to ensure all transitive
+                # dependencies are also not marked as "metadata only"
+                is_metadata = False
+                if hasattr(obj, c.SQLMESH_METADATA):
+                    delattr(obj, c.SQLMESH_METADATA)
+
             if hasattr(obj, c.SQLMESH_MACRO):
                 # We only need to add the undecorated code of @macro() functions in env, which
                 # is accessible through the `__wrapped__` attribute added by functools.wraps
@@ -308,6 +320,11 @@ def walk(obj: t.Any, name: str) -> None:
                 or not hasattr(obj_module, "__file__")
                 or not _is_relative_to(obj_module.__file__, path)
             ):
+                if is_metadata:
+                    setattr(obj, c.SQLMESH_METADATA, True)
+                elif hasattr(obj, c.SQLMESH_METADATA):
+                    delattr(obj, c.SQLMESH_METADATA)
+
                 env[name] = obj
                 return
         elif env[name] != obj:
@@ -318,10 +335,10 @@ def walk(obj: t.Any, name: str) -> None:
         if inspect.isclass(obj):
             for var in decorator_vars(obj):
                 if obj_module and var in obj_module.__dict__:
-                    walk(obj_module.__dict__[var], var)
+                    walk(obj_module.__dict__[var], var, is_metadata)
 
             for base in obj.__bases__:
-                walk(base, base.__qualname__)
+                walk(base, base.__qualname__, is_metadata)
 
             for k, v in obj.__dict__.items():
                 if k.startswith("__"):
@@ -335,19 +352,23 @@ def walk(obj: t.Any, name: str) -> None:
                     # Walk the method if it's part of the object, else it's a global function and we just store it
                     if v.__qualname__.startswith(obj.__qualname__):
                         for k, v in func_globals(v).items():
-                            walk(v, k)
+                            walk(v, k, is_metadata)
                     else:
-                        walk(v, v.__name__)
+                        walk(v, v.__name__, is_metadata)
         elif callable(obj):
             for k, v in func_globals(obj).items():
-                walk(v, k)
+                walk(v, k, is_metadata)
+
+        if is_metadata:
+            setattr(obj, c.SQLMESH_METADATA, True)
 
         # We store the object in the environment after its dependencies, because otherwise we
         # could crash at environment hydration time, since dicts are ordered and the top-level
         # objects would be loaded before their dependencies.
         env[name] = obj
 
-    walk(obj, name)
+    # The "metadata only" annotation of the object is transitive
+    walk(obj, name, getattr(obj, c.SQLMESH_METADATA, None))
 
 
 @dataclass
@@ -411,6 +432,8 @@ def serialize_env(env: t.Dict[str, t.Any], path: Path) -> t.Dict[str, Executable
     serialized = {}
 
     for k, v in env.items():
+        is_metadata = getattr(v, c.SQLMESH_METADATA, None)
+
         if isinstance(v, LITERALS) or v is None:
             serialized[k] = Executable.value(v)
         elif inspect.ismodule(v):
@@ -423,6 +446,7 @@ def serialize_env(env: t.Dict[str, t.Any], path: Path) -> t.Dict[str, Executable
             serialized[k] = Executable(
                 payload=f"import {name}{postfix}",
                 kind=ExecutableKind.IMPORT,
+                is_metadata=is_metadata,
             )
         elif callable(v):
             name = v.__name__
@@ -447,6 +471,7 @@ def serialize_env(env: t.Dict[str, t.Any], path: Path) -> t.Dict[str, Executable
                     v = wrapped
                     file_path = Path(inspect.getfile(wrapped))
                     relative_obj_file_path = _is_relative_to(file_path, path)
+                    is_metadata = is_metadata or getattr(v, c.SQLMESH_METADATA, None)
             except TypeError:
                 file_path = None
                 relative_obj_file_path = False
@@ -459,12 +484,13 @@ 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),
+                    is_metadata=is_metadata,
                 )
             else:
                 serialized[k] = Executable(
                     payload=f"from {v.__module__} import {name}",
                     kind=ExecutableKind.IMPORT,
+                    is_metadata=is_metadata,
                 )
         else:
             raise SQLMeshError(

tests/core/test_model.py

@@ -8665,3 +8665,130 @@ def test_data_hash_unchanged_when_column_type_uses_default_dialect():
 
     # int == int64 in bigquery
     assert model.data_hash == deserialized_model.data_hash
+
+
+def test_transitive_dependency_of_metadata_only_object_is_metadata_only(tmp_path: Path) -> None:
+    init_example_project(tmp_path, dialect="duckdb", template=ProjectTemplate.EMPTY)
+
+    test_model = tmp_path / "models/test_model.sql"
+    test_model.parent.mkdir(parents=True, exist_ok=True)
+    test_model.write_text("MODEL (name test_model, kind FULL); @metadata_macro(); SELECT 1 AS c")
+
+    metadata_macro_code = """
+from sqlglot import parse_one
+from sqlmesh import macro
+
+def get_parsed_query():
+    return parse_one("SELECT 1")
+
+@macro(metadata_only=True)
+def metadata_macro(evaluator):
+    if evaluator.runtime_stage == "evaluating":
+        evaluator.engine_adapter.execute({query})"""
+
+    metadata_macro = tmp_path / "macros/metadata_macro.py"
+    metadata_macro.parent.mkdir(parents=True, exist_ok=True)
+    metadata_macro.write_text(metadata_macro_code.format(query='"SELECT 1"'))
+
+    ctx = Context(
+        config=Config(model_defaults=ModelDefaultsConfig(dialect="duckdb")),
+        paths=tmp_path,
+    )
+
+    model = ctx.get_model("test_model")
+    empty_executable = Executable(payload="")
+
+    assert len(model.python_env) == 1
+    assert (model.python_env.get("metadata_macro") or empty_executable).is_metadata
+
+    ctx.plan(no_prompts=True, auto_apply=True)
+
+    # This should make `parse_one` a dependency and so it'll be added in the python env
+    metadata_macro.write_text(metadata_macro_code.format(query='parse_one("SELECT 1")'))
+
+    ctx.load()
+    model = ctx.get_model("test_model")
+
+    assert len(model.python_env) == 2
+    assert (model.python_env.get("metadata_macro") or empty_executable).is_metadata
+    assert (model.python_env.get("parse_one") or empty_executable).is_metadata
+
+    plan = ctx.plan(no_prompts=True, auto_apply=True)
+    ctx_diff = plan.context_diff
+
+    assert len(ctx_diff.modified_snapshots) == 1
+
+    new_snapshot, _ = ctx_diff.modified_snapshots['"test_model"']
+    assert new_snapshot.change_category == SnapshotChangeCategory.METADATA
+
+    # This should make `get_parsed_query` a dependency and so it'll be added in
+    # the python env, carrying `parse_one` with it as another transitive dependency
+    metadata_macro.write_text(metadata_macro_code.format(query="get_parsed_query()"))
+
+    ctx.load()
+    model = ctx.get_model("test_model")
+
+    assert len(model.python_env) == 3
+    assert (model.python_env.get("metadata_macro") or empty_executable).is_metadata
+    assert (model.python_env.get("get_parsed_query") or empty_executable).is_metadata
+    assert (model.python_env.get("parse_one") or empty_executable).is_metadata
+
+    plan = ctx.plan(no_prompts=True, auto_apply=True)
+    ctx_diff = plan.context_diff
+
+    assert len(ctx_diff.modified_snapshots) == 1
+
+    new_snapshot, _ = ctx_diff.modified_snapshots['"test_model"']
+    assert new_snapshot.change_category == SnapshotChangeCategory.METADATA
+
+
+def test_non_metadata_object_takes_precedence_over_metadata_only_object(tmp_path: Path) -> None:
+    init_example_project(tmp_path, dialect="duckdb", template=ProjectTemplate.EMPTY)
+
+    test_model = tmp_path / "models/test_model.sql"
+    test_model.parent.mkdir(parents=True, exist_ok=True)
+    test_model.write_text("MODEL (name test_model, kind FULL); @m1(); @m2(); SELECT 1 AS c")
+
+    macro_code = """
+from sqlglot import parse_one
+from sqlmesh import macro
+
+def common_dep():
+    pass
+
+def m1_dep():
+    pass
+
+@macro(metadata_only=True)
+def m1(evaluator):
+    m1_dep()
+    common_dep()
+    if evaluator.runtime_stage == "evaluating":
+        evaluator.engine_adapter.execute(parse_one("SELECT 1"))
+
+@macro()
+def m2(evaluator):
+    common_dep()
+    if evaluator.runtime_stage == "evaluating":
+        evaluator.engine_adapter.execute(parse_one("SELECT 1"))"""
+
+    test_macros = tmp_path / "macros/test_macros.py"
+    test_macros.parent.mkdir(parents=True, exist_ok=True)
+    test_macros.write_text(macro_code)
+
+    ctx = Context(
+        config=Config(model_defaults=ModelDefaultsConfig(dialect="duckdb")),
+        paths=tmp_path,
+    )
+    model = ctx.get_model("test_model")
+    empty_executable = Executable(payload="")
+
+    # Both `m1` and `m2` refer to `parse_one`, so for `m1` it would be a transitive metadata-only
+    # object, but since the python env is a flat namespace and `parse_one` is also a dependency
+    # of `m2`, it needs to be treated as non-metadata.
+    assert len(model.python_env) == 5
+    assert (model.python_env.get("m1") or empty_executable).is_metadata
+    assert (model.python_env.get("m1_dep") or empty_executable).is_metadata
+    assert not (model.python_env.get("m2") or empty_executable).is_metadata
+    assert not (model.python_env.get("parse_one") or empty_executable).is_metadata
+    assert not (model.python_env.get("common_dep") or empty_executable).is_metadata