Fix!: mark vars referenced in metadata macros as metadata

Author: georgesittas

sqlmesh/core/constants.py

@@ -80,7 +80,9 @@
 DEFAULT_SCHEMA = "default"
 
 SQLMESH_VARS = "__sqlmesh__vars__"
+SQLMESH_VARS_METADATA = "__sqlmesh__vars__metadata__"
 SQLMESH_BLUEPRINT_VARS = "__sqlmesh__blueprint__vars__"
+SQLMESH_BLUEPRINT_VARS_METADATA = "__sqlmesh__blueprint__vars__metadata__"
 
 VAR = "var"
 BLUEPRINT_VAR = "blueprint_var"

sqlmesh/core/macros.py

@@ -210,7 +210,12 @@ def __init__(
                 self.macros[normalize_macro_name(k)] = self.env[k]
             elif v.is_value:
                 value = self.env[k]
-                if k in (c.SQLMESH_VARS, c.SQLMESH_BLUEPRINT_VARS):
+                if k in (
+                    c.SQLMESH_VARS,
+                    c.SQLMESH_VARS_METADATA,
+                    c.SQLMESH_BLUEPRINT_VARS,
+                    c.SQLMESH_BLUEPRINT_VARS_METADATA,
+                ):
                     value = {
                         var_name: (
                             self.parse_one(var_value.sql)
@@ -528,17 +533,25 @@ def views(self) -> t.List[str]:
 
     def var(self, var_name: str, default: t.Optional[t.Any] = None) -> t.Optional[t.Any]:
         """Returns the value of the specified variable, or the default value if it doesn't exist."""
-        return (self.locals.get(c.SQLMESH_VARS) or {}).get(var_name.lower(), default)
+        return (
+            self.locals.get(c.SQLMESH_VARS) or self.locals.get(c.SQLMESH_VARS_METADATA) or {}
+        ).get(var_name.lower(), default)
 
     def blueprint_var(self, var_name: str, default: t.Optional[t.Any] = None) -> t.Optional[t.Any]:
         """Returns the value of the specified blueprint variable, or the default value if it doesn't exist."""
-        return (self.locals.get(c.SQLMESH_BLUEPRINT_VARS) or {}).get(var_name.lower(), default)
+        return (
+            self.locals.get(c.SQLMESH_BLUEPRINT_VARS)
+            or self.locals.get(c.SQLMESH_BLUEPRINT_VARS_METADATA)
+            or {}
+        ).get(var_name.lower(), default)
 
     @property
     def variables(self) -> t.Dict[str, t.Any]:
         return {
             **self.locals.get(c.SQLMESH_VARS, {}),
+            **self.locals.get(c.SQLMESH_VARS_METADATA, {}),
             **self.locals.get(c.SQLMESH_BLUEPRINT_VARS, {}),
+            **self.locals.get(c.SQLMESH_BLUEPRINT_VARS_METADATA, {}),
         }
 
     def _coerce(self, expr: exp.Expression, typ: t.Any, strict: bool = False) -> t.Any:

sqlmesh/core/model/common.py

@@ -5,6 +5,7 @@
 from pathlib import Path
 
 from astor import to_source
+from collections import defaultdict
 from difflib import get_close_matches
 from sqlglot import exp
 from sqlglot.helper import ensure_list
@@ -28,7 +29,7 @@
     from sqlmesh.utils import registry_decorator
     from sqlmesh.utils.jinja import MacroReference
 
-    MacroCallable = registry_decorator
+    MacroCallable = t.Union[Executable, registry_decorator]
 
 
 def make_python_env(
@@ -48,13 +49,17 @@ def make_python_env(
     dialect: DialectType = None,
 ) -> t.Dict[str, Executable]:
     python_env = {} if python_env is None else python_env
-    variables = variables or {}
     env: t.Dict[str, t.Tuple[t.Any, t.Optional[bool]]] = {}
-    used_macros: t.Dict[
-        str,
-        t.Tuple[t.Union[Executable | MacroCallable], t.Optional[bool]],
-    ] = {}
-    used_variables = (used_variables or set()).copy()
+
+    variables = variables or {}
+    blueprint_variables = blueprint_variables or {}
+
+    used_macros: t.Dict[str, t.Tuple[MacroCallable, t.Optional[bool]]] = {}
+    used_variable_referenced_in_metadata_expression = dict.fromkeys(used_variables or set(), False)
+
+    # 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)
 
     expressions = ensure_list(expressions)
     for expression_metadata in expressions:
@@ -77,16 +82,27 @@ def make_python_env(
                     macros[name],
                     used_macros.get(name, (None, is_metadata))[1] and is_metadata,
                 )
-                if name == c.VAR:
+                if name in (c.VAR, c.BLUEPRINT_VAR):
                     args = macro_func_or_var.this.expressions
                     if len(args) < 1:
-                        raise_config_error("Macro VAR requires at least one argument", path)
+                        raise_config_error(
+                            f"Macro {name.upper()} requires at least one argument", path
+                        )
+
                     if not args[0].is_string:
                         raise_config_error(
                             f"The variable name must be a string literal, '{args[0].sql()}' was given instead",
                             path,
                         )
-                    used_variables.add(args[0].this.lower())
+
+                    var_name = args[0].this.lower()
+                    used_variable_referenced_in_metadata_expression[var_name] = (
+                        used_variable_referenced_in_metadata_expression.get(var_name, True)
+                        and bool(is_metadata)
+                    )
+                else:
+                    for var_ref in _extract_macro_func_variable_references(macro_func_or_var):
+                        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:
@@ -95,17 +111,23 @@ def make_python_env(
                         macros[name],
                         used_macros.get(name, (None, is_metadata))[1] and is_metadata,
                     )
-                elif name in variables:
-                    used_variables.add(name)
+                elif name in variables or name in blueprint_variables:
+                    used_variable_referenced_in_metadata_expression[name] = (
+                        used_variable_referenced_in_metadata_expression.get(name, True)
+                        and bool(is_metadata)
+                    )
             elif (
                 isinstance(macro_func_or_var, (exp.Identifier, d.MacroStrReplace, d.MacroSQL))
             ) and "@" in macro_func_or_var.name:
                 for _, identifier, braced_identifier, _ in MacroStrTemplate.pattern.findall(
                     macro_func_or_var.name
                 ):
                     var_name = braced_identifier or identifier
-                    if var_name in variables:
-                        used_variables.add(var_name)
+                    if var_name in variables or var_name in blueprint_variables:
+                        used_variable_referenced_in_metadata_expression[var_name] = (
+                            used_variable_referenced_in_metadata_expression.get(var_name, True)
+                            and bool(is_metadata)
+                        )
 
     for macro_ref in jinja_macro_references or set():
         if macro_ref.package is None and macro_ref.name in macros:
@@ -126,41 +148,97 @@ def make_python_env(
     python_env.update(serialize_env(env, path=module_path))
     return _add_variables_to_python_env(
         python_env,
-        used_variables,
+        used_variable_referenced_in_metadata_expression,
         variables,
         blueprint_variables=blueprint_variables,
         dialect=dialect,
         strict_resolution=strict_resolution,
+        macro_funcs_by_used_var=macro_funcs_by_used_var,
     )
 
 
+def _extract_macro_func_variable_references(macro_func: exp.Expression) -> t.Set[str]:
+    references = set()
+
+    for n in macro_func.walk():
+        if n is macro_func:
+            continue
+
+        # 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.
+        if isinstance(n, d.MacroFunc):
+            this = n.this
+            if this.name.lower() in (c.VAR, c.BLUEPRINT_VAR) and this.expressions:
+                references.add(this.expressions[0].this.lower())
+        elif isinstance(n, d.MacroVar):
+            references.add(n.name.lower())
+        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)
+            )
+
+    return references
+
+
 def _add_variables_to_python_env(
     python_env: t.Dict[str, Executable],
-    used_variables: t.Optional[t.Set[str]],
+    used_variable_referenced_in_metadata_expression: t.Dict[str, bool],
     variables: t.Optional[t.Dict[str, t.Any]],
     strict_resolution: bool = True,
     blueprint_variables: t.Optional[t.Dict[str, t.Any]] = None,
     dialect: DialectType = None,
+    macro_funcs_by_used_var: t.Optional[t.DefaultDict[str, t.Set[str]]] = None,
 ) -> t.Dict[str, Executable]:
-    _, python_used_variables = parse_dependencies(
+    _, python_used_variable_referenced_in_metadata_expression = parse_dependencies(
         python_env,
         None,
         strict_resolution=strict_resolution,
         variables=variables,
         blueprint_variables=blueprint_variables,
     )
-    used_variables = (used_variables or set()) | python_used_variables
+    for var_name, is_metadata in python_used_variable_referenced_in_metadata_expression.items():
+        used_variable_referenced_in_metadata_expression[var_name] = (
+            used_variable_referenced_in_metadata_expression.get(var_name, True) and is_metadata
+        )
+
+    metadata_used_variables = set()
+    for used_var, macro_names in (macro_funcs_by_used_var or {}).items():
+        if used_variable_referenced_in_metadata_expression.get(used_var) or all(
+            name in python_env and python_env[name].is_metadata for name in macro_names
+        ):
+            metadata_used_variables.add(used_var)
+
+    used_variables = set(used_variable_referenced_in_metadata_expression)
+    non_metadata_used_variables = used_variables - metadata_used_variables
+
+    metadata_variables = {
+        k: v for k, v in (variables or {}).items() if k in metadata_used_variables
+    }
+    variables = {k: v for k, v in (variables or {}).items() if k in non_metadata_used_variables}
 
-    variables = {k: v for k, v in (variables or {}).items() if k in used_variables}
     if variables:
         python_env[c.SQLMESH_VARS] = Executable.value(variables)
+    if metadata_variables:
+        python_env[c.SQLMESH_VARS_METADATA] = Executable.value(metadata_variables, is_metadata=True)
 
     if blueprint_variables:
+        metadata_blueprint_variables = {
+            k: SqlValue(sql=v.sql(dialect=dialect)) if isinstance(v, exp.Expression) else v
+            for k, v in blueprint_variables.items()
+            if k in metadata_used_variables
+        }
         blueprint_variables = {
             k: SqlValue(sql=v.sql(dialect=dialect)) if isinstance(v, exp.Expression) else v
             for k, v in blueprint_variables.items()
+            if k in non_metadata_used_variables
         }
-        python_env[c.SQLMESH_BLUEPRINT_VARS] = Executable.value(blueprint_variables)
+        if blueprint_variables:
+            python_env[c.SQLMESH_BLUEPRINT_VARS] = Executable.value(blueprint_variables)
+        if metadata_blueprint_variables:
+            python_env[c.SQLMESH_BLUEPRINT_VARS_METADATA] = Executable.value(
+                blueprint_variables, is_metadata=True
+            )
 
     return python_env
 
@@ -171,7 +249,7 @@ def parse_dependencies(
     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]]:
+) -> t.Tuple[t.Set[str], t.Dict[str, bool]]:
     """
     Parses the source of a model function and finds upstream table dependencies
     and referenced variables based on calls to context / evaluator.
@@ -185,7 +263,8 @@ def parse_dependencies(
         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.
+        A tuple containing the set of upstream table dependencies and a mapping of
+        the referenced variables associated with their metadata status.
     """
 
     class VariableResolutionContext:
@@ -203,12 +282,16 @@ def blueprint_var(var_name: str, default: t.Optional[t.Any] = None) -> t.Optiona
     local_env = dict.fromkeys(("context", "evaluator"), VariableResolutionContext)
 
     depends_on = set()
-    used_variables = set()
+    used_variable_referenced_in_metadata_expression: t.Dict[str, bool] = {}
 
     for executable in python_env.values():
         if not executable.is_definition:
             continue
+
+        is_metadata = executable.is_metadata
         for node in ast.walk(ast.parse(executable.payload)):
+            next_variables = set()
+
             if isinstance(node, ast.Call):
                 func = node.func
                 if not isinstance(func, ast.Attribute) or not isinstance(func.value, ast.Name):
@@ -239,26 +322,35 @@ 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:
-                    used_variables.add(get_first_arg("var_name").lower())
+                elif func.value.id in ("context", "evaluator") and func.attr in (
+                    c.VAR,
+                    c.BLUEPRINT_VAR,
+                ):
+                    next_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.
-                used_variables.add(c.GATEWAY)
+                next_variables.add(c.GATEWAY)
             elif isinstance(node, ast.FunctionDef) and node.name == entrypoint:
-                used_variables.update(
+                next_variables.update(
                     [
                         arg.arg
                         for arg in [*node.args.args, *node.args.kwonlyargs]
                         if arg.arg != "context"
                     ]
                 )
 
-    return depends_on, used_variables
+            for var_name in next_variables:
+                used_variable_referenced_in_metadata_expression[var_name] = (
+                    used_variable_referenced_in_metadata_expression.get(var_name, True)
+                    and bool(is_metadata)
+                )
+
+    return depends_on, used_variable_referenced_in_metadata_expression
 
 
 def validate_extra_and_required_fields(

sqlmesh/core/model/definition.py

@@ -1774,12 +1774,17 @@ def render(
         start, end = make_inclusive(start or c.EPOCH, end or c.EPOCH, self.dialect)
         execution_time = to_datetime(execution_time or c.EPOCH)
 
-        variables = env.get(c.SQLMESH_VARS, {})
-        variables.update(kwargs.pop("variables", {}))
-
+        variables = {
+            **env.get(c.SQLMESH_VARS, {}),
+            **env.get(c.SQLMESH_VARS_METADATA, {}),
+            **kwargs.pop("variables", {}),
+        }
         blueprint_variables = {
             k: d.parse_one(v.sql, dialect=self.dialect) if isinstance(v, SqlValue) else v
-            for k, v in env.get(c.SQLMESH_BLUEPRINT_VARS, {}).items()
+            for k, v in {
+                **env.get(c.SQLMESH_BLUEPRINT_VARS, {}),
+                **env.get(c.SQLMESH_BLUEPRINT_VARS_METADATA, {}),
+            }.items()
         }
         try:
             kwargs = {

sqlmesh/core/renderer.py

@@ -231,6 +231,7 @@ def _resolve_table(table: str | exp.Table) -> str:
 
         if variables:
             macro_evaluator.locals.setdefault(c.SQLMESH_VARS, {}).update(variables)
+            macro_evaluator.locals.setdefault(c.SQLMESH_VARS_METADATA, {})
 
         for definition in self._macro_definitions:
             try:

sqlmesh/utils/jinja.py

@@ -229,7 +229,7 @@ def extract_macro_references_and_variables(
                     )
 
         for call_name, node in extract_call_names(jinja_str):
-            if call_name[0] == c.VAR:
+            if call_name[0] in (c.VAR, c.BLUEPRINT_VAR):
                 assert isinstance(node, nodes.Call)
                 args = [jinja_call_arg_name(arg) for arg in node.args]
                 if args and args[0]: