Repository: TobikoData/sqlmesh
Pull request #2898 (2024-07-12 20:33:20)

Fix: override python env's variables with unit test vars when present


Patch:
diff --git a/sqlmesh/core/model/definition.py b/sqlmesh/core/model/definition.py
index 0f25d40a..3ff30036 100644
--- a/sqlmesh/core/model/definition.py
+++ b/sqlmesh/core/model/definition.py
@@ -1431,7 +1431,10 @@ class PythonModel(_Model):
         env = prepare_env(self.python_env)
         start, end = make_inclusive(start or c.EPOCH, end or c.EPOCH)
         execution_time = to_datetime(execution_time or c.EPOCH)
+
         variables = env.get(c.SQLMESH_VARS, {})
+        variables.update(kwargs.pop("variables", {}))
+
         try:
             kwargs = {
                 **variables,
diff --git a/sqlmesh/core/renderer.py b/sqlmesh/core/renderer.py
index 9ef7961c..95d35288 100644
--- a/sqlmesh/core/renderer.py
+++ b/sqlmesh/core/renderer.py
@@ -125,8 +125,9 @@ class BaseExpressionRenderer:
             **kwargs,
         }
 
+        variables = kwargs.pop("variables", {})
         jinja_env = self._jinja_macro_registry.build_environment(
-            **{**render_kwargs, **prepare_env(self._python_env)},
+            **{**render_kwargs, **prepare_env(self._python_env), **variables},
             snapshots=(snapshots or {}),
             table_mapping=table_mapping,
             deployability_index=deployability_index,
@@ -179,6 +180,9 @@ class BaseExpressionRenderer:
 
         macro_evaluator.locals.update(render_kwargs)
 
+        if variables:
+            macro_evaluator.locals.setdefault(c.SQLMESH_VARS, {}).update(variables)
+
         resolved_expressions: t.List[t.Optional[exp.Expression]] = []
 
         for expression in expressions:


Test patch:
diff --git a/sqlmesh/core/test/definition.py b/sqlmesh/core/test/definition.py
index dbd56b2d..0871f5ce 100644
--- a/sqlmesh/core/test/definition.py
+++ b/sqlmesh/core/test/definition.py
@@ -32,6 +32,8 @@ if t.TYPE_CHECKING:
 
     Row = t.Dict[str, t.Any]
 
+TIME_KWARG_KEYS = {"start", "end", "execution_time", "latest"}
+
 
 class ModelTest(unittest.TestCase):
     __test__ = False
@@ -545,14 +547,19 @@ class SqlModelTest(ModelTest):
             self.assert_equal(expected, actual, sort=sort, partial=partial)
 
     def _render_model_query(self) -> exp.Query:
-        return self.model.render_query_or_raise(
-            **self.body.get("vars", {}),
+        variables = self.body.get("vars", {}).copy()
+        time_kwargs = {key: variables.pop(key, None) for key in TIME_KWARG_KEYS}
+
+        query = self.model.render_query_or_raise(
+            **time_kwargs,
+            variables=variables,
             engine_adapter=self.engine_adapter,
             table_mapping={
                 name: self._test_fixture_table(name).sql() for name in self.body.get("inputs", {})
             },
             runtime_stage=RuntimeStage.TESTING,
         )
+        return query
 
 
 class PythonModelTest(ModelTest):
@@ -618,7 +625,9 @@ class PythonModelTest(ModelTest):
         time_ctx = freeze_time(self._execution_time) if self._execution_time else nullcontext()
         with patch.dict(self._test_adapter_dialect.generator_class.TRANSFORMS, self._transforms):
             with t.cast(AbstractContextManager, time_ctx):
-                df = next(self.model.render(context=self.context, **self.body.get("vars", {})))
+                variables = self.body.get("vars", {}).copy()
+                time_kwargs = {key: variables.pop(key, None) for key in TIME_KWARG_KEYS}
+                df = next(self.model.render(context=self.context, **time_kwargs, **variables))
                 assert not isinstance(df, exp.Expression)
                 return df if isinstance(df, pd.DataFrame) else df.toPandas()
 
diff --git a/tests/core/test_test.py b/tests/core/test_test.py
index 184b2b13..97163f4c 100644
--- a/tests/core/test_test.py
+++ b/tests/core/test_test.py
@@ -1459,10 +1459,16 @@ def test_variable_usage(tmp_path: Path) -> None:
     )
     context = Context(paths=tmp_path, config=config)
 
-    parent = _create_model("SELECT 1 as id", meta="MODEL (name silver_db.sch.b)")
+    parent = _create_model(
+        "SELECT 1 AS id, '2022-01-01 01:00:00'::TIMESTAMP AS ts",
+        meta="MODEL (name silver_db.sch.b)",
+    )
     parent = t.cast(SqlModel, context.upsert_model(parent))
 
-    child = _create_model("SELECT id FROM silver_db.sch.b", meta="MODEL (name gold_db.sch.a)")
+    child = _create_model(
+        "SELECT @IF(@VAR('myvar'), id, id + 1) AS id FROM silver_db.sch.b",
+        meta="MODEL (name gold_db.sch.a)",
+    )
     child = t.cast(SqlModel, context.upsert_model(child))
 
     test_file = tmp_path / "tests" / "test_parameterized_model_names.yaml"
@@ -1470,12 +1476,16 @@ def test_variable_usage(tmp_path: Path) -> None:
         """
 test_parameterized_model_names:
   model: {{ var('gold') }}.sch.a
+  vars:
+    myvar: True
   inputs:
     {{ var('silver') }}.sch.b:
       - id: 1
+      - id: 2
   outputs:
     query:
       - id: 1
+      - id: 2
         """
     )