Feat: resolve (blueprint) variables when parsing python deps (#4399)

Author: georgesittas

docs/concepts/models/python_models.md

@@ -241,18 +241,16 @@ def execute(
     context.resolve_table("docs_example.another_dependency")
 ```
 
-User-defined [global variables](global-variables) can also be used in `resolve_table` calls, as long as the `depends_on` keyword argument is present and contains the required dependencies. This is shown in the following example:
+User-defined [global variables](global-variables) or [blueprint variables](#python-model-blueprinting) can also be used in `resolve_table` calls, as shown in the following example (similarly for `blueprint_var()`):
 
 ```python linenums="1"
 @model(
     "@schema_name.test_model2",
     kind="FULL",
     columns={"id": "INT"},
-    depends_on=["@schema_name.test_model1"],
 )
 def execute(context, **kwargs):
-    schema_name = context.var("schema_name")
-    table = context.resolve_table(f"{schema_name}.test_model1")
+    table = context.resolve_table(f"{context.var('schema_name')}.test_model1")
     select_query = exp.select("*").from_(table)
     return context.fetchdf(select_query)
 ```

sqlmesh/core/model/common.py

@@ -145,6 +145,8 @@ def _add_variables_to_python_env(
         python_env,
         None,
         strict_resolution=strict_resolution,
+        variables=variables,
+        blueprint_variables=blueprint_variables,
     )
     used_variables = (used_variables or set()) | python_used_variables
 
@@ -163,7 +165,11 @@ def _add_variables_to_python_env(
 
 
 def parse_dependencies(
-    python_env: t.Dict[str, Executable], entrypoint: t.Optional[str], strict_resolution: bool = True
+    python_env: t.Dict[str, Executable],
+    entrypoint: t.Optional[str],
+    strict_resolution: bool = True,
+    variables: t.Optional[t.Dict[str, t.Any]] = None,
+    blueprint_variables: t.Optional[t.Dict[str, t.Any]] = None,
 ) -> t.Tuple[t.Set[str], t.Set[str]]:
     """
     Parses the source of a model function and finds upstream table dependencies
@@ -174,13 +180,29 @@ def parse_dependencies(
         entrypoint: The name of the function.
         strict_resolution: If true, the arguments of `table` and `resolve_table` calls must
             be resolvable at parse time, otherwise an exception will be raised.
+        variables: The variables available to the python environment.
+        blueprint_variables: The blueprint variables available to the python environment.
 
     Returns:
         A tuple containing the set of upstream table dependencies and the set of referenced variables.
     """
+
+    class VariableResolutionContext:
+        """This enables calls like `resolve_table` to reference `var()` and `blueprint_var()`."""
+
+        @staticmethod
+        def var(var_name: str, default: t.Optional[t.Any] = None) -> t.Optional[t.Any]:
+            return (variables or {}).get(var_name.lower(), default)
+
+        @staticmethod
+        def blueprint_var(var_name: str, default: t.Optional[t.Any] = None) -> t.Optional[t.Any]:
+            return (blueprint_variables or {}).get(var_name.lower(), default)
+
     env = prepare_env(python_env)
+    local_env = dict.fromkeys(("context", "evaluator"), VariableResolutionContext)
+
     depends_on = set()
-    variables = set()
+    used_variables = set()
 
     for executable in python_env.values():
         if not executable.is_definition:
@@ -206,7 +228,7 @@ def get_first_arg(keyword_arg_name: str) -> t.Any:
 
                     try:
                         expression = to_source(first_arg)
-                        return eval(expression, env)
+                        return eval(expression, env, local_env)
                     except Exception:
                         if strict_resolution:
                             raise ConfigError(
@@ -217,25 +239,25 @@ def get_first_arg(keyword_arg_name: str) -> t.Any:
                 if func.value.id == "context" and func.attr in ("table", "resolve_table"):
                     depends_on.add(get_first_arg("model_name"))
                 elif func.value.id in ("context", "evaluator") and func.attr == c.VAR:
-                    variables.add(get_first_arg("var_name").lower())
+                    used_variables.add(get_first_arg("var_name").lower())
             elif (
                 isinstance(node, ast.Attribute)
                 and isinstance(node.value, ast.Name)
                 and node.value.id in ("context", "evaluator")
                 and node.attr == c.GATEWAY
             ):
                 # Check whether the gateway attribute is referenced.
-                variables.add(c.GATEWAY)
+                used_variables.add(c.GATEWAY)
             elif isinstance(node, ast.FunctionDef) and node.name == entrypoint:
-                variables.update(
+                used_variables.update(
                     [
                         arg.arg
                         for arg in [*node.args.args, *node.args.kwonlyargs]
                         if arg.arg != "context"
                     ]
                 )
 
-    return depends_on, variables
+    return depends_on, used_variables
 
 
 def single_value_or_tuple(values: t.Sequence) -> exp.Identifier | exp.Tuple:

sqlmesh/core/model/definition.py

@@ -2286,7 +2286,13 @@ def create_python_model(
     dependencies_unspecified = depends_on is None
 
     parsed_depends_on, referenced_variables = (
-        parse_dependencies(python_env, entrypoint, strict_resolution=dependencies_unspecified)
+        parse_dependencies(
+            python_env,
+            entrypoint,
+            strict_resolution=dependencies_unspecified,
+            variables=variables,
+            blueprint_variables=blueprint_variables,
+        )
         if python_env is not None
         else (set(), set())
     )

tests/core/test_model.py

@@ -9861,3 +9861,56 @@ def test_model(context, **kwargs):
     model_executable_str = python_model.render_definition()[1].sql()
     # Make sure the file path is included in the render definition
     assert "# tests/core/test_model.py" in model_executable_str
+
+
+def test_resolve_interpolated_variables_when_parsing_python_deps():
+    @model(
+        name="bla.test_interpolate_var_in_dep_py",
+        kind="full",
+        columns={'"col"': "int"},
+    )
+    def unimportant_testing_model(context, **kwargs):
+        table1 = context.resolve_table(f"{context.var('schema_name')}.table_name")
+        table2 = context.resolve_table(f"{context.blueprint_var('schema_name')}.table_name")
+
+        return context.fetchdf(exp.select("*").from_(table))
+
+    m = model.get_registry()["bla.test_interpolate_var_in_dep_py"].model(
+        module_path=Path("."),
+        path=Path("."),
+        variables={"schema_name": "foo"},
+        blueprint_variables={"schema_name": "baz"},
+    )
+
+    assert m.depends_on == {'"foo"."table_name"', '"baz"."table_name"'}
+    assert m.python_env.get(c.SQLMESH_VARS) == Executable.value({"schema_name": "foo"})
+    assert m.python_env.get(c.SQLMESH_BLUEPRINT_VARS) == Executable.value({"schema_name": "baz"})
+
+    @macro()
+    def unimportant_testing_macro(evaluator, *projections):
+        evaluator.var(f"{evaluator.var('selector')}_variable")
+        evaluator.var(f"{evaluator.blueprint_var('selector')}_variable")
+
+        return exp.select(*[f'{p} AS "{p}"' for p in projections])
+
+    m = load_sql_based_model(
+        d.parse(
+            """
+            MODEL (
+                name bla.test_interpolate_var_in_dep_sql
+            );
+
+            @unimportant_testing_macro();
+
+            SELECT
+              1 AS c
+            """,
+        ),
+        variables={"selector": "bla", "bla_variable": 1, "baz_variable": 2},
+        blueprint_variables={"selector": "baz"},
+    )
+
+    assert m.python_env.get(c.SQLMESH_VARS) == Executable.value(
+        {"selector": "bla", "bla_variable": 1, "baz_variable": 2}
+    )
+    assert m.python_env.get(c.SQLMESH_BLUEPRINT_VARS) == Executable.value({"selector": "baz"})