Fix: Detection of referenced variables in Python SQL models (#2743)

Author: izeigerman

sqlmesh/core/model/definition.py

@@ -1636,6 +1636,8 @@ def create_sql_model(
             used_variables=used_variables,
             path=path,
         )
+    else:
+        python_env = _add_variables_to_python_env(python_env, used_variables, variables)
 
     return _create_model(
         SqlModel,
@@ -1716,6 +1718,8 @@ def create_seed_model(
             used_variables=used_variables,
             path=path,
         )
+    else:
+        python_env = _add_variables_to_python_env(python_env, used_variables, variables)
 
     return _create_model(
         SeedModel,
@@ -1999,15 +2003,22 @@ def _python_env(
             build_env(used_macro.func, env=python_env, name=name, path=module_path)
 
     serialized_env.update(serialize_env(python_env, path=module_path))
+    return _add_variables_to_python_env(serialized_env, used_variables, variables)
+
 
-    _, python_used_variables = _parse_dependencies(serialized_env, None)
-    used_variables |= python_used_variables
+def _add_variables_to_python_env(
+    python_env: t.Dict[str, Executable],
+    used_variables: t.Optional[t.Set[str]],
+    variables: t.Optional[t.Dict[str, t.Any]],
+) -> t.Dict[str, Executable]:
+    _, python_used_variables = _parse_dependencies(python_env, None)
+    used_variables = (used_variables or set()) | python_used_variables
 
     variables = {k: v for k, v in (variables or {}).items() if k in used_variables}
     if variables:
-        serialized_env[c.SQLMESH_VARS] = Executable.value(variables)
+        python_env[c.SQLMESH_VARS] = Executable.value(variables)
 
-    return serialized_env
+    return python_env
 
 
 def _parse_dependencies(

tests/core/test_model.py

@@ -4093,6 +4093,38 @@ def model_with_variables(context, **kwargs):
     assert df.to_dict(orient="records") == [{"a": "test_value", "b": "default_value", "c": None}]
 
 
+def test_variables_python_sql_model(mocker: MockerFixture) -> None:
+    @model(
+        "test_variables_python_model",
+        is_sql=True,
+        kind="full",
+        columns={"a": "string", "b": "string", "c": "string"},
+    )
+    def model_with_variables(evaluator, **kwargs):
+        return exp.select(
+            exp.convert(evaluator.var("TEST_VAR_A")).as_("a"),
+            exp.convert(evaluator.var("test_var_b", "default_value")).as_("b"),
+            exp.convert(evaluator.var("test_var_c")).as_("c"),
+        )
+
+    python_sql_model = model.get_registry()["test_variables_python_model"].model(
+        module_path=Path("."),
+        path=Path("."),
+        variables={"test_var_a": "test_value", "test_var_unused": 2},
+    )
+
+    assert python_sql_model.python_env[c.SQLMESH_VARS] == Executable.value(
+        {"test_var_a": "test_value"}
+    )
+
+    context = ExecutionContext(mocker.Mock(), {}, None, None)
+    query = list(python_sql_model.render(context=context))[0]
+    assert (
+        query.sql()
+        == """SELECT 'test_value' AS "a", 'default_value' AS "b", NULL AS "c" """.strip()
+    )
+
+
 def test_named_variables_python_model(mocker: MockerFixture) -> None:
     @model(
         "test_named_variables_python_model",