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

@@ -183,6 +183,7 @@ def __init__(
         self.default_catalog = default_catalog
         self._path = path
         self._environment_naming_info = environment_naming_info
+        self._blueprint_vars = {}
 
         prepare_env(self.python_env, self.env)
         for k, v in self.python_env.items():
@@ -191,7 +192,13 @@ 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]
+                if k == c.SQLMESH_BLUEPRINT_VARS:
+                    self._blueprint_vars = {
+                        var: self.parse_one(value) for var, value in self.env[k].items()
+                    }
+                    self.locals.update(self._blueprint_vars)
+                else:
+                    self.locals[k] = self.env[k]
 
     def send(
         self, name: str, *args: t.Any, **kwargs: t.Any
@@ -465,7 +472,10 @@ 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.locals.get(c.SQLMESH_VARS) or {}),
+            **self._blueprint_vars,
+        }.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,
@@ -107,6 +114,13 @@ def _add_variables_to_python_env(
     if variables:
         python_env[c.SQLMESH_VARS] = Executable.value(variables)
 
+    if blueprint_variables:
+        blueprint_variables = {
+            k: v.sql(dialect=dialect) if isinstance(v, exp.Expression) else v
+            for k, v in blueprint_variables.items()
+        }
+        python_env[c.SQLMESH_BLUEPRINT_VARS] = Executable.value(blueprint_variables)
+
     return python_env
 
 

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(env.get(c.SQLMESH_BLUEPRINT_VARS, {}))
 
         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,

tests/core/test_model.py

@@ -8200,6 +8200,7 @@ def identity(evaluator, value):
 )
 def entrypoint(context, *args, **kwargs):
     x_var = context.var("x")
+    assert kwargs.get("blueprint")
     return pd.DataFrame({"x": [x_var]})"""
     )
     blueprint_pysql = tmp_path / "models" / "blueprint_sql.py"
@@ -8218,6 +8219,7 @@ def entrypoint(context, *args, **kwargs):
 )
 def entrypoint(evaluator):
     x_var = evaluator.var("x")
+    assert evaluator.var("blueprint")
     return f'SELECT {x_var} AS x'"""
     )
 
@@ -8236,6 +8238,9 @@ def entrypoint(evaluator):
             assert model is not None
             assert "blueprints" not in model.all_fields()
             assert model.python_env.get(c.SQLMESH_VARS) == Executable.value({"x": gateway_no})
+            assert model.python_env.get(c.SQLMESH_BLUEPRINT_VARS) == Executable.value(
+                {"blueprint": f"gw{gateway_no}"}
+            )
             assert context.fetchdf(f"from {model.fqn}").to_dict() == {"x": {0: gateway_no}}
 
     multi_variable_blueprint_example = tmp_path / "models" / "multi_variable_blueprint_example.sql"
@@ -8245,15 +8250,17 @@ def entrypoint(evaluator):
         MODEL (
           name @{customer}.my_table,
           blueprints (
-            (customer := customer1, foo := 'bar'),
-            (customer := customer2, foo := qux),
+            (customer := customer1, customer_field := 'bar'),
+            (customer := customer2, customer_field := qux),
           ),
           kind FULL
         );
 
         SELECT
-          @VAR('foo') AS foo,
-        FROM @VAR('customer').my_source
+          @customer_field AS foo,
+          @{customer_field} AS foo2,
+          @VAR('customer_field') AS foo3,
+        FROM @{customer}.my_source
         """
     )
 
@@ -8266,21 +8273,25 @@ def entrypoint(evaluator):
     customer1_model = models.get('"db"."customer1"."my_table"')
 
     assert customer1_model is not None
-    assert customer1_model.python_env.get(c.SQLMESH_VARS) == Executable.value(
-        {"customer": "customer1", "foo": "'bar'"}
+
+    assert customer1_model.python_env.get(c.SQLMESH_VARS) is None
+    assert customer1_model.python_env.get(c.SQLMESH_BLUEPRINT_VARS) == Executable.value(
+        {"customer": "customer1", "customer_field": "'bar'"}
     )
     assert t.cast(exp.Expression, customer1_model.render_query()).sql() == (
-        """SELECT '''bar''' AS "foo" FROM "db"."customer1"."my_source" AS "my_source\""""
+        """SELECT 'bar' AS "foo", "bar" AS "foo2", 'bar' AS "foo3" FROM "db"."customer1"."my_source" AS "my_source\""""
     )
 
     customer2_model = models.get('"db"."customer2"."my_table"')
 
     assert customer2_model is not None
-    assert customer2_model.python_env.get(c.SQLMESH_VARS) == Executable.value(
-        {"customer": "customer2", "foo": "qux"}
+
+    assert customer2_model.python_env.get(c.SQLMESH_VARS) is None
+    assert customer2_model.python_env.get(c.SQLMESH_BLUEPRINT_VARS) == Executable.value(
+        {"customer": "customer2", "customer_field": "qux"}
     )
     assert t.cast(exp.Expression, customer2_model.render_query()).sql() == (
-        '''SELECT 'qux' AS "foo" FROM "db"."customer2"."my_source" AS "my_source"'''
+        '''SELECT "qux" AS "foo", "qux" AS "foo2", "qux" AS "foo3" FROM "db"."customer2"."my_source" AS "my_source"'''
     )
 
 
@@ -8327,6 +8338,39 @@ def gen_blueprints(evaluator):
     assert '"memory"."customer2"."some_table"' in ctx.models
 
 
+def test_blueprinting_with_quotes(tmp_path: Path) -> None:
+    init_example_project(tmp_path, dialect="duckdb", template=ProjectTemplate.EMPTY)
+
+    template_with_quoted_vars = tmp_path / "models/template_with_quoted_vars.sql"
+    template_with_quoted_vars.parent.mkdir(parents=True, exist_ok=True)
+    template_with_quoted_vars.write_text(
+        """
+        MODEL (
+          name m.@{bp_var},
+          blueprints (
+            (bp_var := "a b"),
+            (bp_var := 'c d'),
+          ),
+        );
+
+        SELECT @bp_var AS c1, @{bp_var} AS c2
+        """
+    )
+
+    ctx = Context(
+        config=Config(model_defaults=ModelDefaultsConfig(dialect="duckdb")), paths=tmp_path
+    )
+    assert len(ctx.models) == 2
+
+    m1 = ctx.get_model('m."a b"', raise_if_missing=True)
+    m2 = ctx.get_model('m."c d"', raise_if_missing=True)
+
+    assert m1.name == 'm."a b"'
+    assert m2.name == 'm."c d"'
+    assert t.cast(exp.Query, m1.render_query()).sql() == '''SELECT "a b" AS "c1", "a b" AS "c2"'''
+    assert t.cast(exp.Query, m2.render_query()).sql() == '''SELECT 'c d' AS "c1", "c d" AS "c2"'''
+
+
 @time_machine.travel("2020-01-01 00:00:00 UTC")
 def test_dynamic_date_spine_model(assert_exp_eq):
     @macro()