Refactor: address new bug, improve testing coverage

georgesittas · georgesittas · commit 142dbc243c41 · 2025-07-30T16:02:28.000+03:00
diff --git a/sqlmesh/core/model/common.py b/sqlmesh/core/model/common.py
@@ -55,12 +55,47 @@ def make_python_env(
     blueprint_variables = blueprint_variables or {}
 
     used_macros: t.Dict[str, t.Tuple[MacroCallable, bool]] = {}
-    used_variables = dict.fromkeys(referenced_variables or set(), False)  # var -> is_metadata
+
+    # var -> True: var is metadata-only
+    # var -> False: var is not metadata-only
+    # var -> None: cannot determine whether var is metadata-only yet, need to walk macros first
+    used_variables: t.Dict[str, t.Optional[bool]] = dict.fromkeys(
+        referenced_variables or set(), False
+    )
+
+    # id(expr) -> true: expr appears under the AST of a metadata-only macro function
+    # id(expr) -> false: expr appears under the AST of a macro function whose metadata status we don't yet know
+    expr_under_metadata_macro_func: t.Dict[int, bool] = {}
 
     # For an expression like @foo(@v1, @bar(@v1, @v2), @v3), the following mapping would be:
     # v1 -> {"foo", "bar"}, v2 -> {"bar"}, v3 -> "foo"
     macro_funcs_by_used_var: t.DefaultDict[str, t.Set[str]] = defaultdict(set)
 
+    def _is_metadata_var(
+        name: str, expression: exp.Expression, appears_in_metadata_expression: bool
+    ) -> t.Optional[bool]:
+        is_metadata_so_far = used_variables.get(name, True)
+        if is_metadata_so_far is False:
+            return False
+
+        appears_under_metadata_macro_func = expr_under_metadata_macro_func.get(id(expression))
+        if is_metadata_so_far and (
+            appears_in_metadata_expression or appears_under_metadata_macro_func
+        ):
+            return True
+
+        if appears_under_metadata_macro_func is False:
+            return None
+
+        return False
+
+    def _is_metadata_macro(name: str, appears_in_metadata_expression: bool) -> bool:
+        if name in used_macros:
+            is_metadata_so_far = used_macros[name][1]
+            return is_metadata_so_far and appears_in_metadata_expression
+
+        return appears_in_metadata_expression
+
     expressions = ensure_list(expressions)
     for expression_metadata in expressions:
         if isinstance(expression_metadata, tuple):
@@ -77,11 +112,8 @@ def make_python_env(
                 if name not in macros:
                     continue
 
-                # If this macro has been seen before as a non-metadata macro, prioritize that
-                used_macros[name] = (
-                    macros[name],
-                    used_macros.get(name, (None, is_metadata))[1] and is_metadata,
-                )
+                used_macros[name] = (macros[name], _is_metadata_macro(name, is_metadata))
+
                 if name in (c.VAR, c.BLUEPRINT_VAR):
                     args = macro_func_or_var.this.expressions
                     if len(args) < 1:
@@ -96,20 +128,22 @@ def make_python_env(
                         )
 
                     var_name = args[0].this.lower()
-                    used_variables[var_name] = used_variables.get(var_name, True) and is_metadata
+                    used_variables[var_name] = _is_metadata_var(
+                        name, macro_func_or_var, is_metadata
+                    )
                 else:
-                    for var_ref in _extract_macro_func_variable_references(macro_func_or_var):
+                    var_refs, _expr_under_metadata_macro_func = (
+                        _extract_macro_func_variable_references(macro_func_or_var, is_metadata)
+                    )
+                    expr_under_metadata_macro_func.update(_expr_under_metadata_macro_func)
+                    for var_ref in var_refs:
                         macro_funcs_by_used_var[var_ref].add(name)
             elif macro_func_or_var.__class__ is d.MacroVar:
                 name = macro_func_or_var.name.lower()
                 if name in macros:
-                    # If this macro has been seen before as a non-metadata macro, prioritize that
-                    used_macros[name] = (
-                        macros[name],
-                        used_macros.get(name, (None, is_metadata))[1] and is_metadata,
-                    )
+                    used_macros[name] = (macros[name], _is_metadata_macro(name, is_metadata))
                 elif name in variables or name in blueprint_variables:
-                    used_variables[name] = used_variables.get(name, True) and is_metadata
+                    used_variables[name] = _is_metadata_var(name, macro_func_or_var, is_metadata)
             elif (
                 isinstance(macro_func_or_var, (exp.Identifier, d.MacroStrReplace, d.MacroSQL))
             ) and "@" in macro_func_or_var.name:
@@ -118,8 +152,8 @@ def make_python_env(
                 ):
                     var_name = braced_identifier or identifier
                     if var_name in variables or var_name in blueprint_variables:
-                        used_variables[var_name] = (
-                            used_variables.get(var_name, True) and is_metadata
+                        used_variables[var_name] = _is_metadata_var(
+                            var_name, macro_func_or_var, is_metadata
                         )
 
     for macro_ref in jinja_macro_references or set():
@@ -150,8 +184,12 @@ def make_python_env(
     )
 
 
-def _extract_macro_func_variable_references(macro_func: exp.Expression) -> t.Set[str]:
+def _extract_macro_func_variable_references(
+    macro_func: exp.Expression,
+    is_metadata: bool,
+) -> t.Tuple[t.Set[str], t.Dict[int, bool]]:
     references = set()
+    expr_under_metadata_macro_func = {}
 
     # Don't descend into nested MacroFunc nodes besides @VAR() and @BLUEPRINT_VAR(), because
     # they will be handled in a separate call of _extract_macro_func_variable_references.
@@ -169,20 +207,23 @@ def _prune_nested_macro_func(expression: exp.Expression) -> bool:
 
             if this.name.lower() in (c.VAR, c.BLUEPRINT_VAR) and args and args[0].is_string:
                 references.add(args[0].this.lower())
+                expr_under_metadata_macro_func[id(n)] = is_metadata
         elif isinstance(n, d.MacroVar):
             references.add(n.name.lower())
+            expr_under_metadata_macro_func[id(n)] = is_metadata
         elif isinstance(n, (exp.Identifier, d.MacroStrReplace, d.MacroSQL)) and "@" in n.name:
             references.update(
                 (braced_identifier or identifier).lower()
                 for _, identifier, braced_identifier, _ in MacroStrTemplate.pattern.findall(n.name)
             )
+            expr_under_metadata_macro_func[id(n)] = is_metadata
 
-    return references
+    return (references, expr_under_metadata_macro_func)
 
 
 def _add_variables_to_python_env(
     python_env: t.Dict[str, Executable],
-    used_variables: t.Dict[str, bool],
+    used_variables: t.Dict[str, t.Optional[bool]],
     variables: t.Optional[t.Dict[str, t.Any]],
     strict_resolution: bool = True,
     blueprint_variables: t.Optional[t.Dict[str, t.Any]] = None,
@@ -197,14 +238,18 @@ def _add_variables_to_python_env(
         blueprint_variables=blueprint_variables,
     )
     for var_name, is_metadata in python_used_variables.items():
-        used_variables[var_name] = used_variables.get(var_name, True) and is_metadata
+        used_variables[var_name] = is_metadata and used_variables.get(var_name)
 
     # Variables are treated as metadata when:
     # - They are only referenced in metadata-only contexts, such as `audits (...)`, virtual statements, etc
     # - They are only referenced in metadata-only macros, either as their arguments or within their definitions
     metadata_used_variables = set()
     for used_var, macro_names in (macro_funcs_by_used_var or {}).items():
-        if used_variables.get(used_var) or all(
+        used_var_is_metadata = used_variables.get(used_var)
+        if used_var_is_metadata is False:
+            continue
+
+        if used_var_is_metadata or all(
             name in python_env and python_env[name].is_metadata for name in macro_names
         ):
             metadata_used_variables.add(used_var)
diff --git a/tests/core/test_model.py b/tests/core/test_model.py
@@ -10050,30 +10050,37 @@ def test_vars_are_taken_into_account_when_propagating_metadata_status(tmp_path:
     test_model.parent.mkdir(parents=True, exist_ok=True)
     test_model.write_text(
         "MODEL (name test_model, kind FULL, blueprints ((v4 := 4, v5 := 5)));"
-        "@m1_with_var();"  # metadata macro, references v1 internally => v1 metadata
-        "@m2_without_var(@v2, @v3);"  # metadata macro => v2 metadata, v3 metadata
-        "@m3_without_var(@v3);"  # non-metadata macro, references v4 => v3, v4 are not metadata
+        "@m1_metadata_references_v1();"  # metadata macro, references v1 internally => v1 metadata
+        "@m2_metadata_does_not_reference_var(@v2, @v3);"  # metadata macro => v2 metadata, v3 metadata
+        "@m3_non_metadata_references_v4(@v3);"  # non-metadata macro, references v4 => v3, v4 are not metadata
         "SELECT 1 AS c;"
+        "@m2_metadata_does_not_reference_var(@v6);"  # metadata macro => v6 is metadata
+        "@m4_non_metadata_references_v6();"  # non-metadata macro, references v6 => v6 is not metadata
         "ON_VIRTUAL_UPDATE_BEGIN;"
-        "@m3_without_var(@v5);"  # non-metadata macro, metadata context => v5 metadata
+        "@m3_non_metadata_references_v4(@v5);"  # non-metadata macro, metadata expression => v5 metadata
         "ON_VIRTUAL_UPDATE_END;"
     )
 
     macro_code = """
 from sqlmesh import macro
 
 @macro(metadata_only=True)
-def m1_with_var(evaluator):
+def m1_metadata_references_v1(evaluator):
     evaluator.var("v1")
     return None
 
 @macro(metadata_only=True)
-def m2_without_var(evaluator, *args):
+def m2_metadata_does_not_reference_var(evaluator, *args):
     return None
 
 @macro()
-def m3_without_var(evaluator, *args):
+def m3_non_metadata_references_v4(evaluator, *args):
     evaluator.var("v4")
+    return None
+
+@macro()
+def m4_non_metadata_references_v6(evaluator):
+    evaluator.var("v6")
     return None"""
 
     test_macros = tmp_path / "macros/test_macros.py"
@@ -10083,23 +10090,24 @@ def m3_without_var(evaluator, *args):
     ctx = Context(
         config=Config(
             model_defaults=ModelDefaultsConfig(dialect="duckdb"),
-            variables={"v1": 1, "v2": 2, "v3": 3},
+            variables={"v1": 1, "v2": 2, "v3": 3, "v6": 6},
         ),
         paths=tmp_path,
     )
     model = ctx.get_model("test_model")
 
     python_env = model.python_env
 
-    assert len(python_env) == 7
-    assert "m1_with_var" in python_env
-    assert "m2_without_var" in python_env
-    assert "m3_without_var" in python_env
+    assert len(python_env) == 8
+    assert "m1_metadata_references_v1" in python_env
+    assert "m2_metadata_does_not_reference_var" in python_env
+    assert "m3_non_metadata_references_v4" in python_env
+    assert "m4_non_metadata_references_v6" in python_env
 
     variables = python_env.get(c.SQLMESH_VARS)
     metadata_variables = python_env.get(c.SQLMESH_VARS_METADATA)
 
-    assert variables == Executable.value({"v1": 1, "v3": 3})
+    assert variables == Executable.value({"v1": 1, "v3": 3, "v6": 6})
     assert metadata_variables == Executable.value({"v2": 2}, is_metadata=True)
 
     blueprint_variables = python_env.get(c.SQLMESH_BLUEPRINT_VARS)
@@ -10115,28 +10123,68 @@ def m3_without_var(evaluator, *args):
     assert macro_evaluator.locals == {
         "runtime_stage": "loading",
         "default_catalog": None,
-        c.SQLMESH_VARS: {"v1": 1, "v3": 3},
+        c.SQLMESH_VARS: {"v1": 1, "v3": 3, "v6": 6},
         c.SQLMESH_VARS_METADATA: {"v2": 2},
         c.SQLMESH_BLUEPRINT_VARS: {"v4": exp.Literal.number("4")},
         c.SQLMESH_BLUEPRINT_VARS_METADATA: {"v5": exp.Literal.number("5")},
     }
     assert macro_evaluator.var("v1") == 1
     assert macro_evaluator.var("v2") == 2
     assert macro_evaluator.var("v3") == 3
+    assert macro_evaluator.var("v6") == 6
     assert macro_evaluator.blueprint_var("v4") == exp.Literal.number("4")
     assert macro_evaluator.blueprint_var("v5") == exp.Literal.number("5")
 
     query_with_vars = macro_evaluator.transform(
-        parse_one("SELECT " + ", ".join(f"@v{var}, @VAR('v{var}')" for var in [1, 2, 3]))
+        parse_one("SELECT " + ", ".join(f"@v{var}, @VAR('v{var}')" for var in [1, 2, 3, 6]))
     )
-    assert t.cast(exp.Expression, query_with_vars).sql() == "SELECT 1, 1, 2, 2, 3, 3"
+    assert t.cast(exp.Expression, query_with_vars).sql() == "SELECT 1, 1, 2, 2, 3, 3, 6, 6"
 
     query_with_blueprint_vars = macro_evaluator.transform(
         parse_one("SELECT " + ", ".join(f"@v{var}, @BLUEPRINT_VAR('v{var}')" for var in [4, 5]))
     )
     assert t.cast(exp.Expression, query_with_blueprint_vars).sql() == "SELECT 4, 4, 5, 5"
 
 
+def test_variable_mentioned_in_both_metadata_and_non_metadata_macro(tmp_path: Path) -> None:
+    init_example_project(tmp_path, engine_type="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_references_v_metadata(); SELECT @m2_references_v_non_metadata() AS c;"
+    )
+
+    macro_code = """
+from sqlmesh import macro
+
+@macro(metadata_only=True)
+def m1_references_v_metadata(evaluator):
+    evaluator.var("v")
+    return None
+
+@macro()
+def m2_references_v_non_metadata(evaluator):
+    evaluator.var("v")
+    return None"""
+
+    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"), variables={"v": 1}),
+        paths=tmp_path,
+    )
+    model = ctx.get_model("test_model")
+
+    python_env = model.python_env
+
+    assert len(python_env) == 3
+    assert set(python_env) > {"m1_references_v_metadata", "m2_references_v_non_metadata"}
+    assert python_env.get(c.SQLMESH_VARS) == Executable.value({"v": 1})
+
+
 def test_non_metadata_object_takes_precedence_over_metadata_only_object(tmp_path: Path) -> None:
     init_example_project(tmp_path, engine_type="duckdb", template=ProjectTemplate.EMPTY)
 
@@ -11167,4 +11215,4 @@ def test_extract_macro_func_variable_references(macro_func: str, variables: t.Se
     from sqlmesh.core.model.common import _extract_macro_func_variable_references
 
     macro_func_ast = parse_one(macro_func)
-    assert _extract_macro_func_variable_references(macro_func_ast) == variables
+    assert _extract_macro_func_variable_references(macro_func_ast, True)[0] == variables
