Fix!: serialize blueprint variables separately to leverage AST

Author: georgesittas

sqlmesh/core/constants.py

@@ -76,6 +76,8 @@
 DEFAULT_SCHEMA = "default"
 
 SQLMESH_VARS = "__sqlmesh__vars__"
+SQLMESH_BLUEPRINT_VARS = "__sqlmesh__blueprint__vars__"
+
 VAR = "var"
 GATEWAY = "gateway"
 

sqlmesh/core/macros.py

@@ -173,14 +173,16 @@ def __init__(
             "MacroEvaluator": MacroEvaluator,
         }
         self.python_env = python_env or {}
-        self._jinja_env: t.Optional[Environment] = jinja_env
         self.macros = {normalize_macro_name(k): v.func for k, v in macro.get_registry().items()}
+        self.columns_to_types_called = False
+        self.default_catalog = default_catalog
+        self.variables: t.Dict[str, t.Any] = {}
+
+        self._jinja_env: t.Optional[Environment] = jinja_env
         self._schema = schema
         self._resolve_table = resolve_table
         self._resolve_tables = resolve_tables
-        self.columns_to_types_called = False
         self._snapshots = snapshots if snapshots is not None else {}
-        self.default_catalog = default_catalog
         self._path = path
         self._environment_naming_info = environment_naming_info
 
@@ -191,7 +193,12 @@ def __init__(
             elif v.is_import and getattr(self.env.get(k), c.SQLMESH_MACRO, None):
                 self.macros[normalize_macro_name(k)] = self.env[k]
             elif v.is_value:
-                self.locals[k] = self.env[k]
+                value = self.env[k]
+                if k == c.SQLMESH_VARS:
+                    self.variables = {**value, **(value.get(c.SQLMESH_BLUEPRINT_VARS) or {})}
+                    self.variables.pop(c.SQLMESH_BLUEPRINT_VARS, None)
+                else:
+                    self.locals[k] = value
 
     def send(
         self, name: str, *args: t.Any, **kwargs: t.Any
@@ -219,7 +226,7 @@ def evaluate_macros(
 
             if isinstance(node, MacroVar):
                 changed = True
-                variables = self.locals.get(c.SQLMESH_VARS, {})
+                variables = self.variables
                 if node.name not in self.locals and node.name.lower() not in variables:
                     if not isinstance(node.parent, StagedFilePath):
                         raise SQLMeshError(f"Macro variable '{node.name}' is undefined.")
@@ -279,17 +286,12 @@ def template(self, text: t.Any, local_variables: t.Dict[str, t.Any]) -> str:
         Returns:
            The rendered string.
         """
-        mapping = {}
-
-        variables = self.locals.get(c.SQLMESH_VARS, {})
-
-        for k, v in chain(variables.items(), self.locals.items(), local_variables.items()):
-            # try to convert all variables into sqlglot expressions
-            # because they're going to be converted into strings in sql
-            # we don't convert strings because that would result in adding quotes
-            if k != c.SQLMESH_VARS:
-                mapping[k] = convert_sql(v, self.dialect)
-
+        # We try to convert all variables into sqlglot expressions because they're going to be converted
+        # into strings; in sql we don't convert strings because that would result in adding quotes
+        mapping = {
+            k: convert_sql(v, self.dialect)
+            for k, v in chain(self.variables.items(), self.locals.items(), local_variables.items())
+        }
         return MacroStrTemplate(str(text)).safe_substitute(mapping)
 
     def evaluate(self, node: MacroFunc) -> exp.Expression | t.List[exp.Expression] | None:
@@ -465,7 +467,7 @@ def gateway(self) -> t.Optional[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.variables.get(var_name.lower(), default)
 
     def _coerce(self, expr: exp.Expression, typ: t.Any, strict: bool = False) -> t.Any:
         """Coerces the given expression to the specified type on a best-effort basis."""

sqlmesh/core/model/common.py

@@ -17,6 +17,7 @@
 from sqlmesh.utils.pydantic import ValidationInfo, field_validator
 
 if t.TYPE_CHECKING:
+    from sqlglot.dialects.dialect import DialectType
     from sqlmesh.utils.jinja import MacroReference
 
 
@@ -30,6 +31,8 @@ def make_python_env(
     path: t.Optional[str | Path] = None,
     python_env: t.Optional[t.Dict[str, Executable]] = None,
     strict_resolution: bool = True,
+    blueprint_variables: t.Optional[t.Dict[str, t.Any]] = None,
+    dialect: DialectType = None,
 ) -> t.Dict[str, Executable]:
     python_env = {} if python_env is None else python_env
     variables = variables or {}
@@ -86,6 +89,8 @@ def make_python_env(
         python_env,
         used_variables,
         variables,
+        blueprint_variables=blueprint_variables,
+        dialect=dialect,
         strict_resolution=strict_resolution,
     )
 
@@ -95,6 +100,8 @@ def _add_variables_to_python_env(
     used_variables: t.Optional[t.Set[str]],
     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,
 ) -> t.Dict[str, Executable]:
     _, python_used_variables = parse_dependencies(
         python_env,
@@ -104,6 +111,20 @@ def _add_variables_to_python_env(
     used_variables = (used_variables or set()) | python_used_variables
 
     variables = {k: v for k, v in (variables or {}).items() if k in used_variables}
+
+    serialized_blueprint_variables = {}
+    for k, v in (blueprint_variables or {}).items():
+        if isinstance(v, exp.Expression):
+            serialized_blueprint_variables[k] = {
+                "value": v.sql(dialect=dialect),
+                "dialect": str(dialect or ""),
+            }
+        else:
+            serialized_blueprint_variables[k] = {"value": v}
+
+    if serialized_blueprint_variables:
+        variables[c.SQLMESH_BLUEPRINT_VARS] = serialized_blueprint_variables
+
     if variables:
         python_env[c.SQLMESH_VARS] = Executable.value(variables)
 

sqlmesh/core/model/decorator.py

@@ -122,6 +122,7 @@ def model(
         default_catalog: t.Optional[str] = None,
         variables: t.Optional[t.Dict[str, t.Any]] = None,
         infer_names: t.Optional[bool] = False,
+        blueprint_variables: t.Optional[t.Dict[str, t.Any]] = None,
     ) -> Model:
         """Get the model registered by this function."""
         env: t.Dict[str, t.Any] = {}
@@ -155,6 +156,7 @@ def model(
             path=path,
             dialect=dialect,
             default_catalog=default_catalog,
+            blueprint_variables=blueprint_variables,
         )
 
         rendered_name = rendered_fields["name"]
@@ -193,6 +195,7 @@ def model(
             "macros": macros,
             "jinja_macros": jinja_macros,
             "audit_definitions": audit_definitions,
+            "blueprint_variables": blueprint_variables,
             **rendered_fields,
         }
 

sqlmesh/core/model/definition.py

@@ -1748,6 +1748,7 @@ def render(
 
         variables = env.get(c.SQLMESH_VARS, {})
         variables.update(kwargs.pop("variables", {}))
+        variables.update(variables.get(c.SQLMESH_BLUEPRINT_VARS) or {})
 
         try:
             kwargs = {
@@ -1855,18 +1856,14 @@ def _extract_blueprints(blueprints: t.Any, path: Path) -> t.List[t.Any]:
     return []  # This is unreachable, but is done to satisfy mypy
 
 
-def _extract_blueprint_variables(
-    blueprint: t.Any,
-    dialect: DialectType,
-    path: Path,
-) -> t.Dict[str, str]:
+def _extract_blueprint_variables(blueprint: t.Any, path: Path) -> t.Dict[str, t.Any]:
     if not blueprint:
         return {}
     if isinstance(blueprint, exp.Paren):
         blueprint = blueprint.unnest()
-        return {blueprint.left.name: blueprint.right.sql(dialect=dialect)}
+        return {blueprint.left.name: blueprint.right}
     if isinstance(blueprint, (exp.Tuple, exp.Array)):
-        return {e.left.name: e.right.sql(dialect=dialect) for e in blueprint.expressions}
+        return {e.left.name: e.right for e in blueprint.expressions}
     if isinstance(blueprint, dict):
         return blueprint
 
@@ -1889,18 +1886,18 @@ def create_models_from_blueprints(
 ) -> t.List[Model]:
     model_blueprints: t.List[Model] = []
     for blueprint in _extract_blueprints(blueprints, path):
-        variables = _extract_blueprint_variables(blueprint, dialect, path)
+        blueprint_variables = _extract_blueprint_variables(blueprint, path)
 
         if gateway:
             rendered_gateway = render_expression(
                 expression=exp.maybe_parse(gateway, dialect=dialect),
                 module_path=module_path,
                 macros=loader_kwargs.get("macros"),
                 jinja_macros=loader_kwargs.get("jinja_macros"),
-                variables=variables,
                 path=path,
                 dialect=dialect,
                 default_catalog=loader_kwargs.get("default_catalog"),
+                blueprint_variables=blueprint_variables,
             )
             gateway_name = rendered_gateway[0].name if rendered_gateway else None
         else:
@@ -1911,7 +1908,8 @@ def create_models_from_blueprints(
                 path=path,
                 module_path=module_path,
                 dialect=dialect,
-                variables={**get_variables(gateway_name), **variables},
+                variables=get_variables(gateway_name),
+                blueprint_variables=blueprint_variables,
                 **loader_kwargs,
             )
         )
@@ -1983,6 +1981,7 @@ def load_sql_based_model(
     default_catalog: t.Optional[str] = None,
     variables: t.Optional[t.Dict[str, t.Any]] = None,
     infer_names: t.Optional[bool] = False,
+    blueprint_variables: t.Optional[t.Dict[str, t.Any]] = None,
     **kwargs: t.Any,
 ) -> Model:
     """Load a model from a parsed SQLMesh model SQL file.
@@ -2059,6 +2058,7 @@ def load_sql_based_model(
         path=path,
         dialect=dialect,
         default_catalog=default_catalog,
+        blueprint_variables=blueprint_variables,
     )
 
     if rendered_meta_exprs is None or len(rendered_meta_exprs) != 1:
@@ -2143,6 +2143,7 @@ def load_sql_based_model(
         variables=variables,
         default_audits=default_audits,
         inline_audits=inline_audits,
+        blueprint_variables=blueprint_variables,
         **meta_fields,
     )
 
@@ -2247,6 +2248,7 @@ def create_python_model(
     module_path: Path = Path(),
     depends_on: t.Optional[t.Set[str]] = None,
     variables: t.Optional[t.Dict[str, t.Any]] = None,
+    blueprint_variables: t.Optional[t.Dict[str, t.Any]] = None,
     **kwargs: t.Any,
 ) -> Model:
     """Creates a Python model.
@@ -2259,6 +2261,7 @@ def create_python_model(
         path: An optional path to the model definition file.
         depends_on: The custom set of model's upstream dependencies.
         variables: The variables to pass to the model.
+        blueprint_variables: The blueprint's variables to pass to the model.
     """
     # Find dependencies for python models by parsing code if they are not explicitly defined
     # Also remove self-references that are found
@@ -2307,6 +2310,7 @@ def create_python_model(
         jinja_macros=jinja_macros,
         module_path=module_path,
         variables=variables,
+        blueprint_variables=blueprint_variables,
         **kwargs,
     )
 
@@ -2361,6 +2365,7 @@ def _create_model(
     macros: t.Optional[MacroRegistry] = None,
     signal_definitions: t.Optional[SignalRegistry] = None,
     variables: t.Optional[t.Dict[str, t.Any]] = None,
+    blueprint_variables: t.Optional[t.Dict[str, t.Any]] = None,
     **kwargs: t.Any,
 ) -> Model:
     _validate_model_fields(klass, {"name", *kwargs} - {"grain", "table_properties"}, path)
@@ -2469,6 +2474,8 @@ def _create_model(
         path=path,
         python_env=python_env,
         strict_resolution=depends_on is None,
+        blueprint_variables=blueprint_variables,
+        dialect=dialect,
     )
 
     env: t.Dict[str, t.Any] = {}
@@ -2632,6 +2639,7 @@ def render_meta_fields(
     dialect: DialectType,
     variables: t.Optional[t.Dict[str, t.Any]],
     default_catalog: t.Optional[str],
+    blueprint_variables: t.Optional[t.Dict[str, t.Any]] = None,
 ) -> t.Dict[str, t.Any]:
     def render_field_value(value: t.Any) -> t.Any:
         if isinstance(value, exp.Expression) or (isinstance(value, str) and "@" in value):
@@ -2645,6 +2653,7 @@ def render_field_value(value: t.Any) -> t.Any:
                 path=path,
                 dialect=dialect,
                 default_catalog=default_catalog,
+                blueprint_variables=blueprint_variables,
             )
             if not rendered_expr:
                 raise SQLMeshError(
@@ -2752,6 +2761,7 @@ def render_expression(
     dialect: DialectType = None,
     variables: t.Optional[t.Dict[str, t.Any]] = None,
     default_catalog: t.Optional[str] = None,
+    blueprint_variables: t.Optional[t.Dict[str, t.Any]] = None,
 ) -> t.Optional[t.List[exp.Expression]]:
     meta_python_env = make_python_env(
         expressions=expression,
@@ -2760,6 +2770,7 @@ def render_expression(
         macros=macros or macro.get_registry(),
         variables=variables,
         path=path,
+        blueprint_variables=blueprint_variables,
     )
     return ExpressionRenderer(
         expression,

sqlmesh/core/renderer.py

@@ -221,7 +221,7 @@ def _render(
         macro_evaluator.locals.update(render_kwargs)
 
         if variables:
-            macro_evaluator.locals.setdefault(c.SQLMESH_VARS, {}).update(variables)
+            macro_evaluator.variables.update(variables)
 
         resolved_expressions: t.List[t.Optional[exp.Expression]] = []
 

sqlmesh/utils/metaprogramming.py

@@ -18,6 +18,7 @@
 from astor import to_source
 
 from sqlmesh.core import constants as c
+from sqlmesh.core.dialect import parse_one
 from sqlmesh.utils import format_exception, unique
 from sqlmesh.utils.errors import SQLMeshError
 from sqlmesh.utils.pydantic import PydanticModel
@@ -491,11 +492,33 @@ def prepare_env(
         python_env.items(), key=lambda item: 0 if item[1].is_import else 1
     ):
         if executable.is_value:
-            env[name] = ast.literal_eval(executable.payload)
+            literal = ast.literal_eval(executable.payload)
+            if (
+                isinstance(literal, dict)
+                and c.SQLMESH_VARS == name
+                and c.SQLMESH_BLUEPRINT_VARS in literal
+            ):
+                bueprint_variables = {}
+                for var, value_metadata in literal[c.SQLMESH_BLUEPRINT_VARS].items():
+                    assert isinstance(value_metadata, dict)
+                    assert "value" in value_metadata
+
+                    value = value_metadata["value"]
+                    dialect = value_metadata.get("dialect")
+
+                    # The presence of dialect signals that the original blueprint value was an AST
+                    bueprint_variables[var] = (
+                        value if dialect is None else parse_one(value, dialect=dialect)
+                    )
+
+                literal[c.SQLMESH_BLUEPRINT_VARS] = bueprint_variables
+
+            env[name] = literal
         else:
             exec(executable.payload, env)
             if executable.alias and executable.name:
                 env[executable.alias] = env[executable.name]
+
     return env