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):