Skip to content

Commit 211d93d

Browse files
authored
Generate a fresh EmrContainerOperator request token on each task attempt (#69625)
EmrContainerOperator fixed its client idempotency token once at construction. When the same operator instance was reused across task retries, every attempt resent that token, so EMR on EKS returned the original job run instead of starting a new one. A retry after a transient submit or startup failure therefore replayed the failed run and could never recover. The default token is now resolved on each execute() call, so a task retry starts a genuinely new job run. An explicitly supplied token is still used verbatim, preserving intentional idempotency for callers who want repeated submissions to collapse to one run.
1 parent 9964065 commit 211d93d

2 files changed

Lines changed: 49 additions & 4 deletions

File tree

‎providers/amazon/src/airflow/providers/amazon/aws/operators/emr.py‎

Lines changed: 7 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -471,8 +471,10 @@ class EmrContainerOperator(AwsBaseOperator[EmrContainerHook]):
471471
:param configuration_overrides: The configuration overrides for the job run,
472472
specifically either application configuration or monitoring configuration.
473473
:param client_request_token: The client idempotency token of the job run request.
474-
Use this if you want to specify a unique ID to prevent two jobs from getting started.
475-
If no token is provided, a UUIDv4 token will be generated for you.
474+
Pass an explicit value to make repeated submissions idempotent: EMR on EKS treats a
475+
resubmission with the same token as the original run, so task retries return that run
476+
instead of starting a new one. If no token is provided, a fresh UUIDv4 is generated on
477+
every task attempt, so each retry starts a genuinely new job run.
476478
:param aws_conn_id: The Airflow connection used for AWS credentials.
477479
If this is ``None`` or empty then the default boto3 behaviour is used. If
478480
running Airflow in a distributed manner and aws_conn_id is None or
@@ -545,7 +547,7 @@ def __init__(
545547
self.release_label = release_label
546548
self.job_driver = job_driver
547549
self.configuration_overrides = configuration_overrides or {}
548-
self.client_request_token = client_request_token or str(uuid4())
550+
self.client_request_token = client_request_token
549551
self.wait_for_completion = wait_for_completion
550552
self.poll_interval = poll_interval
551553
self.max_polling_attempts = max_polling_attempts
@@ -575,13 +577,14 @@ def execute(self, context: Context) -> str | None:
575577
configuration_overrides, context
576578
)
577579

580+
client_request_token = self.client_request_token or str(uuid4())
578581
self.job_id = self.hook.submit_job(
579582
self.name,
580583
self.execution_role_arn,
581584
self.release_label,
582585
self.job_driver,
583586
configuration_overrides,
584-
self.client_request_token,
587+
client_request_token,
585588
self.tags,
586589
self.job_retry_max_attempts,
587590
)

‎providers/amazon/tests/unit/amazon/aws/operators/test_emr_containers.py‎

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616
# under the License.
1717
from __future__ import annotations
1818

19+
import uuid
1920
from unittest import mock
2021
from unittest.mock import patch
2122

@@ -172,6 +173,47 @@ def test_execute_complete_raises_on_error_event(self):
172173
with pytest.raises(AirflowException, match="Error while running job"):
173174
self.emr_container.execute_complete(context=None, event=event)
174175

176+
def _make_operator_without_token(self):
177+
return EmrContainerOperator(
178+
task_id="start_job",
179+
name="test_emr_job",
180+
virtual_cluster_id="vzw123456",
181+
execution_role_arn="arn:aws:somerole",
182+
release_label="6.3.0-latest",
183+
job_driver={},
184+
configuration_overrides={},
185+
poll_interval=0,
186+
)
187+
188+
def test_default_client_request_token_not_generated_at_construction(self):
189+
operator = self._make_operator_without_token()
190+
assert operator.client_request_token is None
191+
192+
@mock.patch.object(EmrContainerHook, "submit_job")
193+
@mock.patch.object(EmrContainerHook, "check_query_status", return_value="COMPLETED")
194+
def test_execute_generates_fresh_token_per_attempt(self, mock_check_query_status, mock_submit_job):
195+
operator = self._make_operator_without_token()
196+
197+
operator.execute(None)
198+
operator.execute(None)
199+
200+
tokens = [call.args[5] for call in mock_submit_job.call_args_list]
201+
assert tokens[0] != tokens[1]
202+
for token in tokens:
203+
assert uuid.UUID(token).version == 4
204+
205+
@mock.patch.object(EmrContainerHook, "submit_job")
206+
@mock.patch.object(EmrContainerHook, "check_query_status", return_value="COMPLETED")
207+
def test_execute_honors_explicit_token_across_attempts(self, mock_check_query_status, mock_submit_job):
208+
self.emr_container.execute(None)
209+
self.emr_container.execute(None)
210+
211+
tokens = [call.args[5] for call in mock_submit_job.call_args_list]
212+
assert tokens == [GENERATED_UUID, GENERATED_UUID]
213+
214+
def test_client_request_token_not_templated(self):
215+
assert "client_request_token" not in EmrContainerOperator.template_fields
216+
175217

176218
class TestEmrEksCreateClusterOperator:
177219
def setup_method(self):

0 commit comments

Comments
 (0)