Fix!: make metadata_only a transitive property in Python objects (#4080)

Author: georgesittas

sqlmesh/core/model/common.py

@@ -24,11 +24,17 @@
 
 if t.TYPE_CHECKING:
     from sqlglot.dialects.dialect import DialectType
+    from sqlmesh.utils import registry_decorator
     from sqlmesh.utils.jinja import MacroReference
 
+    MacroCallable = registry_decorator
+
 
 def make_python_env(
-    expressions: t.Union[exp.Expression, t.List[exp.Expression]],
+    expressions: t.Union[
+        exp.Expression,
+        t.List[t.Union[exp.Expression, t.Tuple[exp.Expression, bool]]],
+    ],
     jinja_macro_references: t.Optional[t.Set[MacroReference]],
     module_path: Path,
     macros: MacroRegistry,
@@ -42,53 +48,79 @@ def make_python_env(
 ) -> t.Dict[str, Executable]:
     python_env = {} if python_env is None else python_env
     variables = variables or {}
-    env: t.Dict[str, t.Any] = {}
-    used_macros = {}
+    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()
 
     expressions = ensure_list(expressions)
-    for expression in expressions:
-        if not isinstance(expression, d.Jinja):
-            for macro_func_or_var in expression.find_all(d.MacroFunc, d.MacroVar, exp.Identifier):
-                if macro_func_or_var.__class__ is d.MacroFunc:
-                    name = macro_func_or_var.this.name.lower()
-                    if name in macros:
-                        used_macros[name] = macros[name]
-                        if name == c.VAR:
-                            args = macro_func_or_var.this.expressions
-                            if len(args) < 1:
-                                raise_config_error("Macro VAR 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())
-                elif macro_func_or_var.__class__ is d.MacroVar:
-                    name = macro_func_or_var.name.lower()
-                    if name in macros:
-                        used_macros[name] = macros[name]
-                    elif name in variables:
-                        used_variables.add(name)
-                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)
+    for expression_metadata in expressions:
+        if isinstance(expression_metadata, tuple):
+            expression, is_metadata = expression_metadata
+        else:
+            expression, is_metadata = expression_metadata, None
+
+        if isinstance(expression, d.Jinja):
+            continue
+
+        for macro_func_or_var in expression.find_all(d.MacroFunc, d.MacroVar, exp.Identifier):
+            if macro_func_or_var.__class__ is d.MacroFunc:
+                name = macro_func_or_var.this.name.lower()
+                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],
+                )
+                if name == c.VAR:
+                    args = macro_func_or_var.this.expressions
+                    if len(args) < 1:
+                        raise_config_error("Macro VAR 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())
+            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],
+                    )
+                elif name in variables:
+                    used_variables.add(name)
+            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)
 
     for macro_ref in jinja_macro_references or set():
         if macro_ref.package is None and macro_ref.name in macros:
-            used_macros[macro_ref.name] = macros[macro_ref.name]
+            used_macros[macro_ref.name] = (macros[macro_ref.name], None)
 
-    for name, used_macro in used_macros.items():
+    for name, (used_macro, is_metadata) in used_macros.items():
         if isinstance(used_macro, Executable):
             python_env[name] = used_macro
         elif not hasattr(used_macro, c.SQLMESH_BUILTIN) and name not in python_env:
-            build_env(used_macro.func, env=env, name=name, path=module_path)
+            build_env(
+                used_macro.func,
+                env=env,
+                name=name,
+                path=module_path,
+                is_metadata_obj=is_metadata,
+            )
 
     python_env.update(serialize_env(env, path=module_path))
     return _add_variables_to_python_env(

sqlmesh/core/model/decorator.py

@@ -125,7 +125,7 @@ def model(
         blueprint_variables: t.Optional[t.Dict[str, t.Any]] = None,
     ) -> Model:
         """Get the model registered by this function."""
-        env: t.Dict[str, t.Any] = {}
+        env: t.Dict[str, t.Tuple[t.Any, t.Optional[bool]]] = {}
         entrypoint = self.func.__name__
 
         if not self.name_provided and not infer_names:

sqlmesh/core/model/definition.py

@@ -2035,14 +2035,7 @@ def load_sql_based_model(
             continue
 
         prop_name = prop.name.lower()
-        if (
-            prop_name
-            in {
-                "signals",
-                "audits",
-            }
-            | PROPERTIES
-        ):
+        if prop_name in {"signals", "audits"} | PROPERTIES:
             unrendered_properties[prop_name] = prop.args.get("value")
         elif (
             prop.name.lower() == "kind"
@@ -2404,14 +2397,17 @@ def _create_model(
         statements.append(kwargs["query"])
     if "post_statements" in kwargs:
         statements.extend(kwargs["post_statements"])
+
+    # Macros extracted from these statements need to be treated as metadata only
     if "on_virtual_update" in kwargs:
-        statements.extend(kwargs["on_virtual_update"])
+        statements.extend((stmt, True) for stmt in kwargs["on_virtual_update"])
 
-    # to allow variables like @gateway to be used in these properties
-    # since rendering shifted from load time to run time
+    # This is done to allow variables like @gateway to be used in these properties
+    # since rendering shifted from load time to run time.
+    # Note: we check for Tuple since that's what we expect from _resolve_properties
     for property_name in PROPERTIES:
-        if property_values := kwargs.get(property_name):
-            statements.extend(property_values)
+        if isinstance(property_values := kwargs.get(property_name), exp.Tuple):
+            statements.extend(property_values.expressions)
 
     jinja_macro_references, used_variables = extract_macro_references_and_variables(
         *(gen(e) for e in statements)
@@ -2443,11 +2439,15 @@ def _create_model(
         raise_config_error(str(ex), location=path)
         raise
 
-    audit_definitions = audit_definitions or {}
-    inline_audits = inline_audits or {}
-    audit_definitions = {**audit_definitions, **inline_audits}
+    audit_definitions = {
+        **(audit_definitions or {}),
+        **(inline_audits or {}),
+    }
 
-    used_audits = set(inline_audits)
+    # TODO: default_audits needs to be merged with model.audits; the former's arguments
+    # are silently dropped today because we add them in audit_definitions. We also need
+    # to check for duplicates when we implement this merging logic.
+    used_audits: t.Set[str] = set()
     used_audits.update(audit_name for audit_name, _ in default_audits or [])
     used_audits.update(audit_name for audit_name, _ in model.audits)
 
@@ -2459,12 +2459,15 @@ def _create_model(
 
     model.audit_definitions.update(audit_definitions)
 
-    statements.extend(audit.query for audit in audit_definitions.values())
+    # Any macro referenced in audits or signals needs to be treated as metadata-only
+    statements.extend((audit.query, True) for audit in audit_definitions.values())
     for _, audit_args in model.audits:
-        statements.extend(audit_args.values())
+        statements.extend(
+            (audit_arg_expression, True) for audit_arg_expression in audit_args.values()
+        )
 
     for _, kwargs in model.signals:
-        statements.extend(kwargs.values())
+        statements.extend((signal_kwarg, True) for signal_kwarg in kwargs.values())
 
     python_env = python_env or {}
 
@@ -2482,7 +2485,7 @@ def _create_model(
         dialect=dialect,
     )
 
-    env: t.Dict[str, t.Any] = {}
+    env: t.Dict[str, t.Tuple[t.Any, t.Optional[bool]]] = {}
 
     for signal_name, _ in model.signals:
         if signal_definitions and signal_name in signal_definitions:

sqlmesh/migrations/v0078_warn_if_non_migratable_python_env.py

@@ -0,0 +1,142 @@
+"""
+This script's goal is to warn users if there is both a metadata and non-metadata reference in
+the python environment of a model. Additionally, it warns them if there's a macro referenced
+in a used audit's query, in the argument list of the audits and signals properties, or in an
+on_virtual_update statement.
+
+Context:
+
+The metadata status for macros and signals is now transitive, i.e. every dependency of a
+metadata macro or signal is also metadata, unless it is referenced by a non-metadata object.
+
+This means that global references of metadata objects may now be excluded from the data hash
+calculation because of their new metadata status, which would lead to a diff.
+
+Additionally, we now implicitly treat macro refs in the aforementioned statements as "metadata-only",
+even though they may not be marked as such by a user. This may also lead to a diff.
+"""
+
+import json
+
+from sqlglot import exp
+
+import sqlmesh.core.dialect as d
+from sqlmesh.core.console import get_console
+
+
+def migrate(state_sync, **kwargs):  # type: ignore
+    engine_adapter = state_sync.engine_adapter
+    schema = state_sync.schema
+    snapshots_table = "_snapshots"
+    if schema:
+        snapshots_table = f"{schema}.{snapshots_table}"
+
+    warning = (
+        "SQLMesh detected that it may not be able to fully migrate the state database. This should not impact "
+        "the migration process, but may result in unexpected changes being reported by the next `sqlmesh plan` "
+        "command. Please run `sqlmesh diff prod` after the migration has completed, before making any new "
+        "changes. If any unexpected changes are reported, consider running a forward-only plan to apply these "
+        "changes and avoid unnecessary backfills: sqlmesh plan prod --forward-only. "
+        "See https://sqlmesh.readthedocs.io/en/stable/concepts/plans/#forward-only-plans for more details.\n"
+    )
+
+    for (snapshot,) in engine_adapter.fetchall(
+        exp.select("snapshot").from_(snapshots_table), quote_identifiers=True
+    ):
+        parsed_snapshot = json.loads(snapshot)
+        node = parsed_snapshot["node"]
+
+        # Standalone audits don't have a data hash, so they're unaffected
+        if node.get("source_type") == "audit":
+            continue
+
+        python_env = node.get("python_env") or {}
+
+        has_metadata = False
+        has_non_metadata = False
+
+        for k, v in python_env.items():
+            if v.get("is_metadata"):
+                has_metadata = True
+            else:
+                has_non_metadata = True
+
+            if has_metadata and has_non_metadata:
+                get_console().log_warning(warning)
+                return
+
+        dialect = node.get("dialect")
+        metadata_hash_statements = []
+
+        # We use try-except here as a conservative measure to avoid any unexpected exceptions
+        try:
+            if on_virtual_update := node.get("on_virtual_update"):
+                metadata_hash_statements.extend(parse_expression(on_virtual_update, dialect))
+
+            for _, audit_args in func_call_validator(node.get("audits") or []):
+                metadata_hash_statements.extend(audit_args.values())
+
+            for signal_name, signal_args in func_call_validator(
+                node.get("signals") or [], is_signal=True
+            ):
+                metadata_hash_statements.extend(signal_args.values())
+
+            if audit_definitions := node.get("audit_definitions"):
+                audit_queries = [
+                    parse_expression(audit["query"], audit["dialect"])
+                    for audit in audit_definitions.values()
+                ]
+                metadata_hash_statements.extend(audit_queries)
+
+            for macro_name in extract_used_macros(metadata_hash_statements):
+                serialized_macro = python_env.get(macro_name)
+                if isinstance(serialized_macro, dict) and not serialized_macro.get("is_metadata"):
+                    get_console().log_warning(warning)
+                    return
+        except Exception:
+            pass
+
+
+def extract_used_macros(expressions):
+    used_macros = set()
+    for expression in expressions:
+        if isinstance(expression, d.Jinja):
+            continue
+
+        for macro_func in expression.find_all(d.MacroFunc):
+            if macro_func.__class__ is d.MacroFunc:
+                used_macros.add(macro_func.this.name.lower())
+
+    return used_macros
+
+
+def func_call_validator(v, is_signal=False):
+    assert isinstance(v, list)
+
+    audits = []
+    for entry in v:
+        if isinstance(entry, dict):
+            args = entry
+            name = "" if is_signal else entry.pop("name")
+        else:
+            assert isinstance(entry, (tuple, list))
+            name, args = entry
+
+        parsed_audit = {
+            key: d.parse_one(value) if isinstance(value, str) else value
+            for key, value in args.items()
+        }
+        audits.append((name.lower(), parsed_audit))
+
+    return audits
+
+
+def parse_expression(v, dialect):
+    if v is None:
+        return None
+
+    if isinstance(v, list):
+        return [d.parse_one(e, dialect=dialect) for e in v]
+
+    assert isinstance(v, str)
+    return d.parse_one(v, dialect=dialect)