Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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,
)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@
# under the License.
from __future__ import annotations

import uuid
from unittest import mock
from unittest.mock import patch

Expand Down Expand Up @@ -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):
Expand Down