Fix teardown TaskFlow tasks failing on reuse and with expand_kwargs

task-sdk/src/airflow/sdk/bases/decorator.py

@@ -561,17 +561,15 @@ def __attrs_post_init__(self):
         update_wrapper(self, self.function)
 
     def __call__(self, *args: FParams.args, **kwargs: FParams.kwargs) -> XComArg:
-        if self.is_teardown:
-            if "trigger_rule" in self.kwargs:
-                raise ValueError("Trigger rule not configurable for teardown tasks.")
-            self.kwargs.update(trigger_rule=TriggerRule.ALL_DONE_SETUP_SUCCESS)
-        on_failure_fail_dagrun = self.kwargs.pop("on_failure_fail_dagrun", self.on_failure_fail_dagrun)
+        task_kwargs = self.kwargs.copy()
+        self._set_teardown_trigger_rule(task_kwargs)
+        on_failure_fail_dagrun = task_kwargs.pop("on_failure_fail_dagrun", self.on_failure_fail_dagrun)
         op = self.operator_class(
             python_callable=self.function,
             op_args=args,
             op_kwargs=kwargs,
             multiple_outputs=self.multiple_outputs,
-            **self.kwargs,
+            **task_kwargs,
         )
         op.is_setup = self.is_setup
         op.is_teardown = self.is_teardown
@@ -587,6 +585,13 @@ def __call__(self, *args: FParams.args, **kwargs: FParams.kwargs) -> XComArg:
             op.doc_md = inspect.cleandoc(self.function.__doc__)
         return XComArg(op)
 
+    def _set_teardown_trigger_rule(self, task_kwargs: dict[str, Any]) -> None:
+        if not self.is_teardown:
+            return
+        if "trigger_rule" in task_kwargs:
+            raise ValueError("Trigger rule not configurable for teardown tasks.")
+        task_kwargs["trigger_rule"] = TriggerRule.ALL_DONE_SETUP_SUCCESS
+
     def _validate_arg_names(self, func: ValidationSource, kwargs: dict[str, Any]):
         # Ensure that context variables are not shadowed.
         task_group = self.kwargs.get("task_group") or TaskGroupContext.get_current(self.kwargs.get("dag"))
@@ -613,10 +618,6 @@ def expand(self, **map_kwargs: OperatorExpandArgument) -> XComArg:
         prevent_duplicates(self.kwargs, map_kwargs, fail_reason="mapping already partial")
         # Since the input is already checked at parse time, we can set strict
         # to False to skip the checks on execution.
-        if self.is_teardown:
-            if "trigger_rule" in self.kwargs:
-                raise ValueError("Trigger rule not configurable for teardown tasks.")
-            self.kwargs.update(trigger_rule=TriggerRule.ALL_DONE_SETUP_SUCCESS)
         return self._expand(DictOfListsExpandInput(map_kwargs), strict=False)
 
     def expand_kwargs(self, kwargs: OperatorExpandKwargsArgument, *, strict: bool = True) -> XComArg:
@@ -647,6 +648,7 @@ def _expand(self, expand_input: ExpandInput, *, strict: bool) -> XComArg:
         ensure_xcomarg_return_value(expand_input.value)
 
         task_kwargs = self.kwargs.copy()
+        self._set_teardown_trigger_rule(task_kwargs)
         dag = task_kwargs.pop("dag", None) or DagContext.get_current()
         task_group = task_kwargs.pop("task_group", None) or TaskGroupContext.get_current(dag)
 

task-sdk/tests/task_sdk/definitions/decorators/test_setup_teardown.py

@@ -20,7 +20,7 @@
 import pytest
 
 from airflow.providers.standard.operators.bash import BashOperator
-from airflow.sdk import DAG, setup, task, task_group, teardown
+from airflow.sdk import DAG, TriggerRule, setup, task, task_group, teardown
 from airflow.sdk.definitions.decorators.setup_teardown import context_wrapper
 from airflow.sdk.exceptions import AirflowException
 
@@ -218,6 +218,62 @@ def mytask():
         assert teardown_task.is_teardown
         assert teardown_task.on_failure_fail_dagrun
 
+    @pytest.mark.parametrize(
+        "first_use",
+        [
+            pytest.param(lambda t: t(), id="call"),
+            pytest.param(lambda t: t.expand(x=[1, 2]), id="expand"),
+        ],
+    )
+    def test_teardown_trigger_rule_set_when_decorated_function_is_reused(self, first_use):
+        @teardown
+        def mytask(x=None):
+            print("I am a teardown task")
+
+        with DAG("test_teardown_trigger_rule_set_when_decorated_function_is_reused") as dag:
+            first_use(mytask)
+            mytask.override(task_id="mytask2")()
+
+        assert len(dag.task_group.children) == 2
+        for teardown_task in dag.task_group.children.values():
+            assert teardown_task.is_teardown
+            assert teardown_task.trigger_rule == TriggerRule.ALL_DONE_SETUP_SUCCESS
+
+    def test_teardown_on_failure_fail_dagrun_override_retained_when_called_twice(self):
+        @teardown
+        def mytask():
+            print("I am a teardown task")
+
+        mytask_fail_dagrun = mytask.override(on_failure_fail_dagrun=True)
+        with DAG("test_teardown_on_failure_fail_dagrun_override_retained_when_called_twice") as dag:
+            mytask_fail_dagrun()
+            mytask_fail_dagrun()
+
+        assert len(dag.task_group.children) == 2
+        for teardown_task in dag.task_group.children.values():
+            assert teardown_task.on_failure_fail_dagrun
+
+    def test_teardown_trigger_rule_set_when_mapped_with_expand_kwargs(self):
+        @teardown
+        def mytask(x):
+            print("I am a teardown task")
+
+        with DAG("test_teardown_trigger_rule_set_when_mapped_with_expand_kwargs") as dag:
+            mytask.expand_kwargs([{"x": 1}, {"x": 2}])
+
+        teardown_task = dag.task_group.children["mytask"]
+        assert teardown_task.is_teardown
+        assert teardown_task.trigger_rule == TriggerRule.ALL_DONE_SETUP_SUCCESS
+
+    def test_teardown_raises_when_trigger_rule_passed_with_expand_kwargs(self):
+        @teardown
+        def mytask(x):
+            print("I am a teardown task")
+
+        with DAG("test_teardown_raises_when_trigger_rule_passed_with_expand_kwargs"):
+            with pytest.raises(ValueError, match="Trigger rule not configurable for teardown tasks."):
+                mytask.override(trigger_rule=TriggerRule.ALL_DONE).expand_kwargs([{"x": 1}])
+
     def test_setup_teardown_mixed_up_in_a_dag(self):
         @setup
         def setuptask():