Repository: TobikoData/sqlmesh
Pull request #2748 (2024-06-06 16:49:30)

Fix: hydrate python env. with vars used in MacroStrReplace, MacroSQL


Patch:
diff --git a/sqlmesh/core/model/definition.py b/sqlmesh/core/model/definition.py
index bffdd89e..aca1bcf7 100644
--- a/sqlmesh/core/model/definition.py
+++ b/sqlmesh/core/model/definition.py
@@ -1983,10 +1983,10 @@ def _python_env(
                     elif name in variables:
                         used_variables.add(name)
                 elif (
-                    isinstance(macro_func_or_var, exp.Identifier) and "@" in macro_func_or_var.this
-                ):
+                    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.this
+                        macro_func_or_var.name
                     ):
                         var_name = braced_identifier or identifier
                         if var_name in variables:
diff --git a/sqlmesh/schedulers/airflow/dag_generator.py b/sqlmesh/schedulers/airflow/dag_generator.py
index c06d9381..9f00c017 100644
--- a/sqlmesh/schedulers/airflow/dag_generator.py
+++ b/sqlmesh/schedulers/airflow/dag_generator.py
@@ -257,6 +257,7 @@ class SnapshotDagGenerator:
 
             finalize_task = self._create_finalize_task(plan_dag_spec.environment)
             before_finalize_task >> finalize_task
+            finalize_task >> end_task
 
             on_plan_apply_end_task = PythonOperator(
                 task_id="on_plan_apply_end",
@@ -266,9 +267,7 @@ class SnapshotDagGenerator:
             )
             finalize_task >> on_plan_apply_end_task
 
-            self._add_notification_target_tasks(
-                plan_dag_spec, start_task, end_task, on_plan_apply_end_task
-            )
+            self._add_notification_target_tasks(plan_dag_spec, start_task, finalize_task)
             return dag
 
     def _add_notification_target_tasks(
@@ -276,9 +275,7 @@ class SnapshotDagGenerator:
         request: common.PlanDagSpec,
         start_task: BaseOperator,
         end_task: BaseOperator,
-        previous_end_task: BaseOperator,
     ) -> None:
-        has_success_or_failed_notification = False
         for notification_target in request.notification_targets:
             notification_operator_provider = NOTIFICATION_TARGET_TO_OPERATOR_PROVIDER.get(
                 type(notification_target)
@@ -297,15 +294,9 @@ class SnapshotDagGenerator:
             if plan_start_notification_task:
                 start_task >> plan_start_notification_task
             if plan_success_notification_task:
-                has_success_or_failed_notification = True
-                previous_end_task >> plan_success_notification_task
-                plan_success_notification_task >> end_task
+                end_task >> plan_success_notification_task
             if plan_failed_notification_task:
-                has_success_or_failed_notification = True
-                previous_end_task >> plan_failed_notification_task
-                plan_failed_notification_task >> end_task
-        if not has_success_or_failed_notification:
-            previous_end_task >> end_task
+                end_task >> plan_failed_notification_task
 
     def _create_creation_tasks(
         self,


Test patch:
diff --git a/tests/core/test_model.py b/tests/core/test_model.py
index 4e655cfc..76baccc3 100644
--- a/tests/core/test_model.py
+++ b/tests/core/test_model.py
@@ -3887,20 +3887,34 @@ def test_variables():
             kind FULL,
         );
 
-        SELECT @VAR('TEST_VAR_A') AS a, @VAR('test_var_b', 'default_value') AS b, @VAR('test_var_c') AS c, @TEST_MACRO_VAR() AS d;
+        SELECT
+          @VAR('TEST_VAR_A') AS a,
+          @VAR('test_var_b', 'default_value') AS b,
+          @VAR('test_var_c') AS c,
+          @TEST_MACRO_VAR() AS d,
+          @'foo_@{test_var_e}' AS e,
+          @SQL(foo_@{test_var_f}) AS f,
+          'foo_@{test_var_unused}' AS g
     """,
         default_dialect="bigquery",
     )
 
     model = load_sql_based_model(
-        expressions, variables={"test_var_a": "test_value", "test_var_d": 1, "test_var_unused": 2}
+        expressions,
+        variables={
+            "test_var_a": "test_value",
+            "test_var_d": 1,
+            "test_var_e": 4,
+            "test_var_f": 5,
+            "test_var_unused": 2,
+        },
     )
     assert model.python_env[c.SQLMESH_VARS] == Executable.value(
-        {"test_var_a": "test_value", "test_var_d": 1}
+        {"test_var_a": "test_value", "test_var_d": 1, "test_var_e": 4, "test_var_f": 5}
     )
     assert (
         model.render_query().sql(dialect="bigquery")
-        == "SELECT 'test_value' AS `a`, 'default_value' AS `b`, NULL AS `c`, 11 AS `d`"
+        == "SELECT 'test_value' AS `a`, 'default_value' AS `b`, NULL AS `c`, 11 AS `d`, 'foo_4' AS `e`, `foo_5` AS `f`, 'foo_@{test_var_unused}' AS `g`"
     )
 
     with pytest.raises(ConfigError, match=r"Macro VAR requires at least one argument.*"):