You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by vi...@apache.org on 2024/06/28 15:49:27 UTC
(airflow) branch main updated: Adding cluster to ecs trigger event to avoid defer error (#40482)
This is an automated email from the ASF dual-hosted git repository.
vincbeck pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 6c12744dd8 Adding cluster to ecs trigger event to avoid defer error (#40482)
6c12744dd8 is described below
commit 6c12744dd8656e1d8b066c7edc8f0ab60ac124d2
Author: ellisms <11...@users.noreply.github.com>
AuthorDate: Fri Jun 28 11:49:20 2024 -0400
Adding cluster to ecs trigger event to avoid defer error (#40482)
* Adding cluster to ecs trigger event to avoid defer error
* restore cluster after defer
---
airflow/providers/amazon/aws/operators/ecs.py | 1 +
airflow/providers/amazon/aws/triggers/ecs.py | 4 +++-
tests/providers/amazon/aws/operators/test_ecs.py | 4 ++--
tests/providers/amazon/aws/triggers/test_ecs.py | 1 +
4 files changed, 7 insertions(+), 3 deletions(-)
diff --git a/airflow/providers/amazon/aws/operators/ecs.py b/airflow/providers/amazon/aws/operators/ecs.py
index 294291dc0a..1cd8685cf2 100644
--- a/airflow/providers/amazon/aws/operators/ecs.py
+++ b/airflow/providers/amazon/aws/operators/ecs.py
@@ -586,6 +586,7 @@ class EcsRunTaskOperator(EcsBaseOperator):
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.
diff --git a/airflow/providers/amazon/aws/triggers/ecs.py b/airflow/providers/amazon/aws/triggers/ecs.py
index 1177aa657a..dd86899f22 100644
--- a/airflow/providers/amazon/aws/triggers/ecs.py
+++ b/airflow/providers/amazon/aws/triggers/ecs.py
@@ -179,7 +179,9 @@ class TaskDoneTrigger(BaseTrigger):
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):
diff --git a/tests/providers/amazon/aws/operators/test_ecs.py b/tests/providers/amazon/aws/operators/test_ecs.py
index 9bccd22f9e..a6915214a0 100644
--- a/tests/providers/amazon/aws/operators/test_ecs.py
+++ b/tests/providers/amazon/aws/operators/test_ecs.py
@@ -674,13 +674,13 @@ class TestEcsRunTaskOperator(EcsBaseTestCase):
@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(
diff --git a/tests/providers/amazon/aws/triggers/test_ecs.py b/tests/providers/amazon/aws/triggers/test_ecs.py
index 27c815be82..a5c11d9298 100644
--- a/tests/providers/amazon/aws/triggers/test_ecs.py
+++ b/tests/providers/amazon/aws/triggers/test_ecs.py
@@ -92,3 +92,4 @@ class TestTaskDoneTrigger:
assert response.payload["status"] == "success"
assert response.payload["task_arn"] == "my_task_arn"
+ assert response.payload["cluster"] == "cluster"