Adding cluster to ecs trigger event to avoid defer error (apache#40482)

* Adding cluster to ecs trigger event to avoid defer error

* restore cluster after defer

airflow/providers/amazon/aws/operators/ecs.py

@@ -586,6 +586,7 @@ def execute_complete(self, context: Context, event: dict[str, Any] | None = None
         if event["status"] != "success":
             raise AirflowException(f"Error in task execution: {event}")
         self.arn = event["task_arn"]  # restore arn to its updated value, needed for next steps
+        self.cluster = event["cluster"]
         self._after_execution()
         if self._aws_logs_enabled():
             # same behavior as non-deferrable mode, return last line of logs of the task.

airflow/providers/amazon/aws/triggers/ecs.py

@@ -179,7 +179,9 @@ async def run(self) -> AsyncIterator[TriggerEvent]:
                         cluster=self.cluster, tasks=[self.task_arn], WaiterConfig={"MaxAttempts": 1}
                     )
                     # we reach this point only if the waiter met a success criteria
-                    yield TriggerEvent({"status": "success", "task_arn": self.task_arn})
+                    yield TriggerEvent(
+                        {"status": "success", "task_arn": self.task_arn, "cluster": self.cluster}
+                    )
                     return
                 except WaiterError as error:
                     if "terminal failure" in str(error):

tests/providers/amazon/aws/operators/test_ecs.py

@@ -674,13 +674,13 @@ def test_with_defer(self, client_mock, xcom_mock):
 
     @mock.patch.object(EcsRunTaskOperator, "client", new_callable=PropertyMock)
     def test_execute_complete(self, client_mock):
-        event = {"status": "success", "task_arn": "my_arn"}
+        event = {"status": "success", "task_arn": "my_arn", "cluster": "test_cluster"}
         self.ecs.reattach = True
 
         self.ecs.execute_complete(None, event)
 
         # task gets described to assert its success
-        client_mock().describe_tasks.assert_called_once_with(cluster="c", tasks=["my_arn"])
+        client_mock().describe_tasks.assert_called_once_with(cluster="test_cluster", tasks=["my_arn"])
 
     @pytest.mark.db_test
     @pytest.mark.parametrize(

tests/providers/amazon/aws/triggers/test_ecs.py

@@ -92,3 +92,4 @@ async def test_run_success(self, _, client_mock):
 
         assert response.payload["status"] == "success"
         assert response.payload["task_arn"] == "my_task_arn"
+        assert response.payload["cluster"] == "cluster"