Fix: 'duplicate definitions' python env bug due to module import order (#4249)

Author: georgesittas

sqlmesh/utils/metaprogramming.py

@@ -61,6 +61,16 @@ def _code_globals(code: types.CodeType) -> t.Dict[str, None]:
     return variables
 
 
+def _globals_match(obj1: t.Any, obj2: t.Any) -> bool:
+    return type(obj1) == type(obj2) and (
+        obj1 == obj2
+        or (
+            getattr(obj1, "__module__", None) == getattr(obj2, "__module__", None)
+            and getattr(obj1, "__name__", None) == getattr(obj2, "__name__", None)
+        )
+    )
+
+
 def func_globals(func: t.Callable) -> t.Dict[str, t.Any]:
     """Finds all global references and closures in a function and nested functions.
 
@@ -287,9 +297,17 @@ def build_env(
 
     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"):
+        if obj_module and obj_module.__name__ == "builtins":
             return
 
+        if name in visited:
+            if name not in env or _globals_match(env[name][0], obj):
+                return
+
+            raise SQLMeshError(
+                f"Cannot store {obj} in environment, duplicate definitions found for '{name}'"
+            )
+
         visited.add(name)
         name_missing_from_env = name not in env
 
@@ -320,7 +338,7 @@ def walk(obj: t.Any, name: str, is_metadata: t.Optional[bool] = None) -> None:
             ):
                 env[name] = (obj, is_metadata)
                 return
-        elif env[name][0] != obj:
+        elif not _globals_match(env[name][0], obj):
             raise SQLMeshError(
                 f"Cannot store {obj} in environment, duplicate definitions found for '{name}'"
             )

tests/core/test_model.py

@@ -9235,3 +9235,129 @@ def test_runtime_stage(evaluator):
     model = load_sql_based_model(expressions, jinja_macros=jinja_macros)
     assert model.render_query().sql() == "SELECT 'loading' AS a, 'loading_bla' AS b"
     assert set(model.python_env) == {"noop", "test_runtime_stage"}
+
+
+def test_python_env_references_are_unequal_but_point_to_same_definition(tmp_path: Path) -> None:
+    # This tests for regressions against an edge case bug which was due to reloading modules
+    # in sqlmesh.utils.metaprogramming.import_python_file. Depending on the module loading
+    # order, we could get a "duplicate symbol in python env" error, even though the references
+    # essentially pointed to the same definition (e.g. function or class).
+    init_example_project(tmp_path, dialect="duckdb", template=ProjectTemplate.EMPTY)
+
+    db_path = str(tmp_path / "db.db")
+    db_connection = DuckDBConnectionConfig(database=db_path)
+    config = Config(
+        gateways={"duckdb": GatewayConfig(connection=db_connection)},
+        model_defaults=ModelDefaultsConfig(dialect="duckdb"),
+    )
+
+    file_a = tmp_path / "macros" / "a.py"
+    file_b = tmp_path / "macros" / "b.py"
+    file_c = tmp_path / "macros" / "c.py"
+
+    file_a.write_text(
+        """from macros.c import target
+
+def f1():
+    target()
+"""
+    )
+    file_b.write_text(
+        """from sqlmesh import macro
+
+from macros.a import f1
+from macros.c import target
+
+@macro()
+def first_macro(evaluator):
+    f1()
+
+@macro()
+def second_macro(evaluator):
+    target()
+"""
+    )
+    file_c.write_text(
+        """def target():
+    pass
+"""
+    )
+
+    model_file = tmp_path / "models" / "model.sql"
+    model_file.write_text("MODEL (name a); @first_macro(); @second_macro(); SELECT 1 AS c")
+
+    ctx = Context(paths=tmp_path, config=config, load=False)
+    loader = ctx._loaders[0]
+
+    original_glob_paths = loader._glob_paths
+
+    def _patched_glob_paths(path, *args, **kwargs):
+        if path == tmp_path / "macros":
+            yield from [file_a, file_c, file_b]
+        else:
+            yield from original_glob_paths(path, *args, **kwargs)
+
+    # We force the import order to be a.py -> c.py -> b.py:
+    #
+    # 1. a.py is loaded, so "macros", "macros.a" and "macros.c" are loaded in sys.modules
+    # 2. c.py is loaded, so "macros" and "macros.c" are reloaded in sys.modules
+    # 3. b.py is loaded, so "macros" is reloaded and "macros.b" is loaded in sys.modules
+    #
+    # (1) => id(sys.modules["macros.a"].target) == id(sys.modules["macros.c"].target) == X
+    # (2) => id(sys.modules["macros.c"].target) == Y != X == id(sys.modules["macros.a"].target)
+    # (3) => affects neither sys.modules["macros.a"] nor sys.modules["macros.c"], just loads the macros
+    #
+    # At this point we have two different function instances, one in sys.modules["macros.a"] and one
+    # in sys.modules["macros.c"], which encapsulate the same definition (source code). This used to
+    # lead to a crash, because we prohibit unequal objects with the same name in the python env.
+    with patch.object(loader, "_glob_paths", side_effect=_patched_glob_paths):
+        ctx.load()
+
+    model = ctx.models['"a"']
+    python_env = model.python_env
+
+    assert len(python_env) == 4
+    assert python_env.get("target") == Executable(
+        payload="def target():\n    pass", name="target", path="macros/c.py"
+    )
+
+
+def test_unequal_duplicate_python_env_references_are_prohibited(tmp_path: Path) -> None:
+    init_example_project(tmp_path, dialect="duckdb", template=ProjectTemplate.EMPTY)
+
+    db_path = str(tmp_path / "db.db")
+    db_connection = DuckDBConnectionConfig(database=db_path)
+    config = Config(
+        gateways={"duckdb": GatewayConfig(connection=db_connection)},
+        model_defaults=ModelDefaultsConfig(dialect="duckdb"),
+    )
+
+    file_a = tmp_path / "macros" / "unimportant_macro.py"
+    file_b = tmp_path / "macros" / "just_f.py"
+
+    file_a.write_text(
+        """from sqlmesh import macro
+from macros.just_f import f
+
+a = False
+
+@macro()
+def unimportant_macro(evaluator):
+    print(a)
+    f()
+    return 1
+"""
+    )
+    file_b.write_text(
+        """a = 0
+
+def f():
+    print(a)
+"""
+    )
+
+    model_file = tmp_path / "models" / "model.sql"
+    model_file.write_text("MODEL (name m); SELECT @unimportant_macro() AS unimportant_macro")
+
+    with pytest.raises(SQLMeshError, match=r"duplicate definitions found"):
+        Context(paths=tmp_path, config=config)