(airflow) branch main updated: Fix EcsRunTaskOperator reattach (#51412)

[el...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

eladkal 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 8abcddd2ebd Fix EcsRunTaskOperator reattach (#51412)
8abcddd2ebd is described below

commit 8abcddd2ebd3ba2d673d4c0380d7cba36b226ca2
Author: Niko Oliveira <on...@amazon.com>
AuthorDate: Fri Jun 6 12:48:04 2025 -0700

    Fix EcsRunTaskOperator reattach (#51412)
    
    Previously the ti primary key was used to create a UUID to uniquely
    identify which task started an ECS task, and if the same task was
    retried the same ECS Task could be used. However, operators now
    receive a RuntimeTaskInstance which does not have an equivalent key, so
    we now construct it with the attributes directly.
---
 providers/amazon/src/airflow/providers/amazon/aws/operators/ecs.py | 2 +-
 providers/amazon/tests/system/amazon/aws/example_ecs.py            | 1 +
 providers/amazon/tests/unit/amazon/aws/operators/test_ecs.py       | 4 ++++
 3 files changed, 6 insertions(+), 1 deletion(-)

diff --git a/providers/amazon/src/airflow/providers/amazon/aws/operators/ecs.py b/providers/amazon/src/airflow/providers/amazon/aws/operators/ecs.py
index c83ae2690e4..7205c396129 100644
--- a/providers/amazon/src/airflow/providers/amazon/aws/operators/ecs.py
+++ b/providers/amazon/src/airflow/providers/amazon/aws/operators/ecs.py
@@ -517,7 +517,7 @@ class EcsRunTaskOperator(EcsBaseOperator):
         if self.reattach:
             # Generate deterministic UUID which refers to unique TaskInstanceKey
             ti: TaskInstance = context["ti"]
-            self._started_by = generate_uuid(*map(str, ti.key.primary))
+            self._started_by = generate_uuid(*map(str, [ti.dag_id, ti.task_id, ti.run_id, ti.map_index]))
             self.log.info("Try to find run with startedBy=%r", self._started_by)
             self._try_reattach_task(started_by=self._started_by)
 
diff --git a/providers/amazon/tests/system/amazon/aws/example_ecs.py b/providers/amazon/tests/system/amazon/aws/example_ecs.py
index 8f427f7674e..5ac6664b21d 100644
--- a/providers/amazon/tests/system/amazon/aws/example_ecs.py
+++ b/providers/amazon/tests/system/amazon/aws/example_ecs.py
@@ -150,6 +150,7 @@ with DAG(
     # [START howto_operator_ecs_run_task]
     run_task = EcsRunTaskOperator(
         task_id="run_task",
+        reattach=True,
         cluster=existing_cluster_name,
         task_definition=register_task.output,
         overrides={
diff --git a/providers/amazon/tests/unit/amazon/aws/operators/test_ecs.py b/providers/amazon/tests/unit/amazon/aws/operators/test_ecs.py
index e8e65119d17..b1987852b88 100644
--- a/providers/amazon/tests/unit/amazon/aws/operators/test_ecs.py
+++ b/providers/amazon/tests/unit/amazon/aws/operators/test_ecs.py
@@ -622,6 +622,10 @@ class TestEcsRunTaskOperator(EcsBaseTestCase):
         """Test reattach on first running Task ARN."""
         mock_ti = mock.MagicMock(name="MockedTaskInstance")
         mock_ti.key.primary = ("mock_dag", "mock_ti", "mock_runid", 42)
+        mock_ti.dag_id = "mock_dag"
+        mock_ti.task_id = "mock_ti"
+        mock_ti.run_id = "mock_runid"
+        mock_ti.map_index = 42
         fake_uuid = "01-02-03-04"
         uuid_mock.return_value = fake_uuid