From be8189674c8a59c80756d20ecf76c67d5d404238 Mon Sep 17 00:00:00 2001 From: Ramit Kataria Date: Wed, 8 Jul 2026 09:26:36 -0700 Subject: [PATCH] Generate a fresh EmrContainerOperator request token on each task attempt 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. --- .../providers/amazon/aws/operators/emr.py | 11 +++-- .../aws/operators/test_emr_containers.py | 42 +++++++++++++++++++ 2 files changed, 49 insertions(+), 4 deletions(-) diff --git a/providers/amazon/src/airflow/providers/amazon/aws/operators/emr.py b/providers/amazon/src/airflow/providers/amazon/aws/operators/emr.py index 14d91d613509b..c1e5ee91556b6 100644 --- a/providers/amazon/src/airflow/providers/amazon/aws/operators/emr.py +++ b/providers/amazon/src/airflow/providers/amazon/aws/operators/emr.py @@ -471,8 +471,10 @@ class EmrContainerOperator(AwsBaseOperator[EmrContainerHook]): :param configuration_overrides: The configuration overrides for the job run, specifically either application configuration or monitoring configuration. :param client_request_token: The client idempotency token of the job run request. - Use this if you want to specify a unique ID to prevent two jobs from getting started. - If no token is provided, a UUIDv4 token will be generated for you. + Pass an explicit value to make repeated submissions idempotent: EMR on EKS treats a + resubmission with the same token as the original run, so task retries return that run + instead of starting a new one. If no token is provided, a fresh UUIDv4 is generated on + every task attempt, so each retry starts a genuinely new job run. :param aws_conn_id: The Airflow connection used for AWS credentials. If this is ``None`` or empty then the default boto3 behaviour is used. If running Airflow in a distributed manner and aws_conn_id is None or @@ -545,7 +547,7 @@ def __init__( self.release_label = release_label self.job_driver = job_driver self.configuration_overrides = configuration_overrides or {} - self.client_request_token = client_request_token or str(uuid4()) + self.client_request_token = client_request_token self.wait_for_completion = wait_for_completion self.poll_interval = poll_interval self.max_polling_attempts = max_polling_attempts @@ -575,13 +577,14 @@ def execute(self, context: Context) -> str | None: configuration_overrides, context ) + client_request_token = self.client_request_token or str(uuid4()) self.job_id = self.hook.submit_job( self.name, self.execution_role_arn, self.release_label, self.job_driver, configuration_overrides, - self.client_request_token, + client_request_token, self.tags, self.job_retry_max_attempts, ) diff --git a/providers/amazon/tests/unit/amazon/aws/operators/test_emr_containers.py b/providers/amazon/tests/unit/amazon/aws/operators/test_emr_containers.py index d18a0f599cb7f..6c01ab14724b0 100644 --- a/providers/amazon/tests/unit/amazon/aws/operators/test_emr_containers.py +++ b/providers/amazon/tests/unit/amazon/aws/operators/test_emr_containers.py @@ -16,6 +16,7 @@ # under the License. from __future__ import annotations +import uuid from unittest import mock from unittest.mock import patch @@ -172,6 +173,47 @@ def test_execute_complete_raises_on_error_event(self): with pytest.raises(AirflowException, match="Error while running job"): self.emr_container.execute_complete(context=None, event=event) + def _make_operator_without_token(self): + return EmrContainerOperator( + task_id="start_job", + name="test_emr_job", + virtual_cluster_id="vzw123456", + execution_role_arn="arn:aws:somerole", + release_label="6.3.0-latest", + job_driver={}, + configuration_overrides={}, + poll_interval=0, + ) + + def test_default_client_request_token_not_generated_at_construction(self): + operator = self._make_operator_without_token() + assert operator.client_request_token is None + + @mock.patch.object(EmrContainerHook, "submit_job") + @mock.patch.object(EmrContainerHook, "check_query_status", return_value="COMPLETED") + def test_execute_generates_fresh_token_per_attempt(self, mock_check_query_status, mock_submit_job): + operator = self._make_operator_without_token() + + operator.execute(None) + operator.execute(None) + + tokens = [call.args[5] for call in mock_submit_job.call_args_list] + assert tokens[0] != tokens[1] + for token in tokens: + assert uuid.UUID(token).version == 4 + + @mock.patch.object(EmrContainerHook, "submit_job") + @mock.patch.object(EmrContainerHook, "check_query_status", return_value="COMPLETED") + def test_execute_honors_explicit_token_across_attempts(self, mock_check_query_status, mock_submit_job): + self.emr_container.execute(None) + self.emr_container.execute(None) + + tokens = [call.args[5] for call in mock_submit_job.call_args_list] + assert tokens == [GENERATED_UUID, GENERATED_UUID] + + def test_client_request_token_not_templated(self): + assert "client_request_token" not in EmrContainerOperator.template_fields + class TestEmrEksCreateClusterOperator: def setup_method(self):