Refactor to track metadata objects using a separate value in env

Author: georgesittas

sqlmesh/core/model/common.py

@@ -42,7 +42,7 @@ def make_python_env(
 ) -> t.Dict[str, Executable]:
     python_env = {} if python_env is None else python_env
     variables = variables or {}
-    env: t.Dict[str, t.Any] = {}
+    env: t.Dict[str, t.Tuple[t.Any, t.Optional[bool]]] = {}
     used_macros = {}
     used_variables = (used_variables or set()).copy()
 

sqlmesh/core/model/decorator.py

@@ -125,7 +125,7 @@ def model(
         blueprint_variables: t.Optional[t.Dict[str, t.Any]] = None,
     ) -> Model:
         """Get the model registered by this function."""
-        env: t.Dict[str, t.Any] = {}
+        env: t.Dict[str, t.Tuple[t.Any, t.Optional[bool]]] = {}
         entrypoint = self.func.__name__
 
         if not self.name_provided and not infer_names:

sqlmesh/core/model/definition.py

@@ -2482,7 +2482,7 @@ def _create_model(
         dialect=dialect,
     )
 
-    env: t.Dict[str, t.Any] = {}
+    env: t.Dict[str, t.Tuple[t.Any, t.Optional[bool]]] = {}
 
     for signal_name, _ in model.signals:
         if signal_definitions and signal_name in signal_definitions:

sqlmesh/utils/metaprogramming.py

@@ -266,7 +266,7 @@ def normalize_source(obj: t.Any) -> str:
 def build_env(
     obj: t.Any,
     *,
-    env: t.Dict[str, t.Any],
+    env: t.Dict[str, t.Tuple[t.Any, t.Optional[bool]]],
     name: str,
     path: Path,
 ) -> None:
@@ -291,16 +291,12 @@ 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, True)):
             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)
+                is_metadata = None
 
             if hasattr(obj, c.SQLMESH_MACRO):
                 # We only need to add the undecorated code of @macro() functions in env, which
@@ -320,14 +316,9 @@ def walk(obj: t.Any, name: str, is_metadata: t.Optional[bool] = None) -> 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
+                env[name] = (obj, is_metadata)
                 return
-        elif env[name] != obj:
+        elif env[name][0] != obj:
             raise SQLMeshError(
                 f"Cannot store {obj} in environment, duplicate definitions found for '{name}'"
             )
@@ -359,13 +350,10 @@ def walk(obj: t.Any, name: str, is_metadata: t.Optional[bool] = None) -> None:
             for k, v in func_globals(obj).items():
                 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
+        env[name] = (obj, is_metadata)
 
     # The "metadata only" annotation of the object is transitive
     walk(obj, name, getattr(obj, c.SQLMESH_METADATA, None))
@@ -416,8 +404,8 @@ 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]:
@@ -431,11 +419,9 @@ 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)
-
+    for k, (v, is_metadata) in env.items():
         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 +457,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

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,21 @@ 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, None)}, path=Path("tests/utils"))
 
     with pytest.raises(SQLMeshError):
-        serialize_env({"select_all": SELECT_ALL}, path=Path("tests/utils"))
+        serialize_env({"select_all": (SELECT_ALL, None)}, path=Path("tests/utils"))
 
 
 def test_serialize_env() -> None:
-    env: t.Dict[str, t.Any] = {}
     path = Path("tests/utils")
+    env: t.Dict[str, t.Tuple[t.Any, t.Optional[bool]]] = {}
+
     build_env(main_func, env=env, name="MAIN", path=path)
-    env = serialize_env(env, path=path)  # type: ignore
+    serialized_env = serialize_env(env, path=path)  # type: ignore
+    assert prepare_env(serialized_env)
 
-    assert prepare_env(env)
-    assert env == {
+    expected_env = {
         "MAIN": Executable(
             name="main_func",
             alias="MAIN",
@@ -244,8 +236,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 +310,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 +385,23 @@ def function_with_custom_decorator():
     return""",
             alias="_func",
         ),
+        "SQLGLOT_META": Executable.value("sqlglot.meta"),
     }
+
+    assert all(is_metadata is None for (_, is_metadata) in env.values())
+    assert serialized_env == expected_env
+
+    # Annotate the entrypoint as "metadata only" to show how it propagates
+    setattr(main_func, c.SQLMESH_METADATA, True)
+
+    env = {}
+
+    build_env(main_func, env=env, name="MAIN", path=path)
+    serialized_env = serialize_env(env, path=path)  # type: ignore
+    assert prepare_env(serialized_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 all(is_metadata for (_, is_metadata) in env.values())
+    assert serialized_env == expected_env