Refactor to track metadata objects using a new dict

georgesittas · georgesittas · commit a9c0f3f72e29 · 2025-04-04T22:39:33.000+03:00
diff --git a/sqlmesh/core/model/common.py b/sqlmesh/core/model/common.py
@@ -43,6 +43,7 @@ def make_python_env(
     python_env = {} if python_env is None else python_env
     variables = variables or {}
     env: t.Dict[str, t.Any] = {}
+    env_metadata: t.Set[str] = set()
     used_macros = {}
     used_variables = (used_variables or set()).copy()
 
@@ -88,9 +89,15 @@ def make_python_env(
         if isinstance(used_macro, Executable):
             python_env[name] = used_macro
         elif not hasattr(used_macro, c.SQLMESH_BUILTIN) and name not in python_env:
-            build_env(used_macro.func, env=env, name=name, path=module_path)
+            build_env(
+                used_macro.func,
+                env=env,
+                env_metadata=env_metadata,
+                name=name,
+                path=module_path,
+            )
 
-    python_env.update(serialize_env(env, path=module_path))
+    python_env.update(serialize_env(env, env_metadata=env_metadata, path=module_path))
     return _add_variables_to_python_env(
         python_env,
         used_variables,
diff --git a/sqlmesh/core/model/decorator.py b/sqlmesh/core/model/decorator.py
@@ -126,6 +126,7 @@ def model(
     ) -> Model:
         """Get the model registered by this function."""
         env: t.Dict[str, t.Any] = {}
+        env_metadata: t.Set[str] = set()
         entrypoint = self.func.__name__
 
         if not self.name_provided and not infer_names:
@@ -145,7 +146,7 @@ def model(
                         f"""Python model "{self.name}"'s `kind` dictionary must contain a `name` key with a valid ModelKindName enum value."""
                     )
 
-        build_env(self.func, env=env, name=entrypoint, path=module_path)
+        build_env(self.func, env=env, env_metadata=env_metadata, name=entrypoint, path=module_path)
 
         rendered_fields = render_meta_fields(
             fields={"name": self.name, **self.kwargs},
@@ -184,7 +185,7 @@ def model(
             "defaults": rendered_defaults,
             "path": path,
             "time_column_format": time_column_format,
-            "python_env": serialize_env(env, path=module_path),
+            "python_env": serialize_env(env, env_metadata=env_metadata, path=module_path),
             "physical_schema_mapping": physical_schema_mapping,
             "project": project,
             "default_catalog": default_catalog,
diff --git a/sqlmesh/core/model/definition.py b/sqlmesh/core/model/definition.py
@@ -2483,15 +2483,16 @@ def _create_model(
     )
 
     env: t.Dict[str, t.Any] = {}
+    env_metadata: t.Set[str] = set()
 
     for signal_name, _ in model.signals:
         if signal_definitions and signal_name in signal_definitions:
             func = signal_definitions[signal_name].func
             setattr(func, c.SQLMESH_METADATA, True)
-            build_env(func, env=env, name=signal_name, path=module_path)
+            build_env(func, env=env, env_metadata=env_metadata, name=signal_name, path=module_path)
 
     model.python_env.update(python_env)
-    model.python_env.update(serialize_env(env, path=module_path))
+    model.python_env.update(serialize_env(env, env_metadata=env_metadata, path=module_path))
     model._path = path
     model.set_time_format(time_column_format)
 
diff --git a/sqlmesh/utils/metaprogramming.py b/sqlmesh/utils/metaprogramming.py
@@ -267,6 +267,7 @@ def build_env(
     obj: t.Any,
     *,
     env: t.Dict[str, t.Any],
+    env_metadata: t.Set[str],
     name: str,
     path: Path,
 ) -> None:
@@ -277,6 +278,7 @@ def build_env(
     Args:
         obj: Any python object.
         env: Dictionary to store the env.
+        env_metadata: Set to store the keys that correspond to "metadata only" objects in the env.
         name: Name of the object in the env.
         path: The module path to serialize. Other modules will not be walked and treated as imports.
     """
@@ -291,16 +293,13 @@ def walk(obj: t.Any, name: str, is_metadata: t.Optional[bool] = None) -> None:
         visited.add(name)
         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 name_missing_from_env or (not is_metadata and env[name] == obj and name in env_metadata):
             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)
+                env_metadata.remove(name)
 
             if hasattr(obj, c.SQLMESH_MACRO):
                 # We only need to add the undecorated code of @macro() functions in env, which
@@ -321,9 +320,7 @@ def walk(obj: t.Any, name: str, is_metadata: t.Optional[bool] = None) -> None:
                 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_metadata.add(name)
 
                 env[name] = obj
                 return
@@ -360,7 +357,7 @@ def walk(obj: t.Any, name: str, is_metadata: t.Optional[bool] = None) -> None:
                 walk(v, k, is_metadata)
 
         if is_metadata:
-            setattr(obj, c.SQLMESH_METADATA, True)
+            env_metadata.add(name)
 
         # 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
@@ -416,26 +413,31 @@ def is_value(self) -> bool:
         return self.kind == ExecutableKind.VALUE
 
     @classmethod
-    def value(cls, v: t.Any) -> Executable:
-        return Executable(payload=repr(v), kind=ExecutableKind.VALUE)
+    def value(cls, v: t.Any, is_metadata: t.Optional[bool] = None) -> Executable:
+        return Executable(payload=repr(v), kind=ExecutableKind.VALUE, is_metadata=is_metadata)
 
 
-def serialize_env(env: t.Dict[str, t.Any], path: Path) -> t.Dict[str, Executable]:
+def serialize_env(
+    env: t.Dict[str, t.Any],
+    env_metadata: t.Set[str],
+    path: Path,
+) -> t.Dict[str, Executable]:
     """Serializes a python function into a self contained dictionary.
 
     Recursively walks a function's globals to store all other references inside of env.
 
     Args:
         env: Dictionary to store the env.
+        env_metadata: Keys that correspond to "metadata only" objects in the env.
         path: The root path to seralize. Other modules will not be walked and treated as imports.
     """
     serialized = {}
 
     for k, v in env.items():
-        is_metadata = getattr(v, c.SQLMESH_METADATA, None)
+        is_metadata = True if k in env_metadata else None
 
         if isinstance(v, LITERALS) or v is None:
-            serialized[k] = Executable.value(v)
+            serialized[k] = Executable.value(v, is_metadata=is_metadata)
         elif inspect.ismodule(v):
             name = v.__name__
             if hasattr(v, "__file__") and _is_relative_to(v.__file__, path):
@@ -471,7 +473,6 @@ 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
diff --git a/tests/utils/test_metaprogramming.py b/tests/utils/test_metaprogramming.py
@@ -11,7 +11,7 @@
 from pytest_mock.plugin import MockerFixture
 from sqlglot import exp
 from sqlglot import exp as expressions
-from sqlglot.expressions import to_table
+from sqlglot.expressions import SQLGLOT_META, to_table
 from sqlglot.optimizer.pushdown_projections import SELECT_ALL
 
 import tests.utils.test_date as test_date
@@ -100,13 +100,6 @@ def other_func(a: int) -> int:
     return X + a + W
 
 
-def noop_metadata() -> None:
-    return None
-
-
-setattr(noop_metadata, c.SQLMESH_METADATA, True)
-
-
 @contextmanager
 def test_context_manager():
     yield
@@ -134,8 +127,7 @@ def main_func(y: int, foo=exp.true(), *, bar=expressions.Literal.number(1) + 2)
     sqlglot.parse_one("1")
     MyClass()
     DataClass(x=y)
-    noop_metadata()
-    normalize_model_name("test")
+    normalize_model_name("test" + SQLGLOT_META)
     fetch_data()
     function_with_custom_decorator()
 
@@ -154,7 +146,6 @@ def test_func_globals() -> None:
         "Z": 3,
         "DataClass": DataClass,
         "MyClass": MyClass,
-        "noop_metadata": noop_metadata,
         "normalize_model_name": normalize_model_name,
         "other_func": other_func,
         "sqlglot": sqlglot,
@@ -163,6 +154,7 @@ def test_func_globals() -> None:
         "fetch_data": fetch_data,
         "test_context_manager": test_context_manager,
         "function_with_custom_decorator": function_with_custom_decorator,
+        "SQLGLOT_META": SQLGLOT_META,
     }
     assert func_globals(other_func) == {
         "X": 1,
@@ -194,8 +186,7 @@ def test_normalize_source() -> None:
     sqlglot.parse_one('1')
     MyClass()
     DataClass(x=y)
-    noop_metadata()
-    normalize_model_name('test')
+    normalize_model_name('test' + SQLGLOT_META)
     fetch_data()
     function_with_custom_decorator()
 
@@ -221,20 +212,23 @@ def closure(z: int):
 def test_serialize_env_error() -> None:
     with pytest.raises(SQLMeshError):
         # pretend to be the module pandas
-        serialize_env({"test_date": test_date}, path=Path("tests/utils"))
+        serialize_env({"test_date": test_date}, set(), path=Path("tests/utils"))
 
     with pytest.raises(SQLMeshError):
-        serialize_env({"select_all": SELECT_ALL}, path=Path("tests/utils"))
+        serialize_env({"select_all": SELECT_ALL}, set(), path=Path("tests/utils"))
 
 
 def test_serialize_env() -> None:
-    env: t.Dict[str, t.Any] = {}
     path = Path("tests/utils")
-    build_env(main_func, env=env, name="MAIN", path=path)
-    env = serialize_env(env, path=path)  # type: ignore
 
+    env: t.Dict[str, t.Any] = {}
+    env_metadata: t.Set[str] = set()
+
+    build_env(main_func, env=env, env_metadata=env_metadata, name="MAIN", path=path)
+    env = serialize_env(env, env_metadata=env_metadata, path=path)  # type: ignore
     assert prepare_env(env)
-    assert env == {
+
+    expected_env = {
         "MAIN": Executable(
             name="main_func",
             alias="MAIN",
@@ -244,8 +238,7 @@ def test_serialize_env() -> None:
     sqlglot.parse_one('1')
     MyClass()
     DataClass(x=y)
-    noop_metadata()
-    normalize_model_name('test')
+    normalize_model_name('test' + SQLGLOT_META)
     fetch_data()
     function_with_custom_decorator()
 
@@ -319,13 +312,6 @@ def test_context_manager():
             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,
-        ),
         "normalize_model_name": Executable(
             payload="from sqlmesh.core.dialect import normalize_model_name",
             kind=ExecutableKind.IMPORT,
@@ -401,4 +387,24 @@ def function_with_custom_decorator():
     return""",
             alias="_func",
         ),
+        "SQLGLOT_META": Executable.value("sqlglot.meta"),
     }
+
+    assert env_metadata == set()
+    assert env == expected_env
+
+    # Annotate the entrypoint as "metadata only" to show how it propagates
+    setattr(main_func, c.SQLMESH_METADATA, True)
+
+    env = {}
+    env_metadata = set()
+
+    build_env(main_func, env=env, env_metadata=env_metadata, name="MAIN", path=path)
+    env = serialize_env(env, env_metadata=env_metadata, path=path)  # type: ignore
+    assert prepare_env(env)
+
+    expected_env = {k: Executable(**v.dict(), is_metadata=True) for k, v in expected_env.items()}
+
+    # Every object is treated as "metadata only", transitively
+    assert env_metadata == set(env)
+    assert env == expected_env
