From 63a9d9f6b5b6df7b836752f518f5469330a5c78c Mon Sep 17 00:00:00 2001 From: Amir Mor Date: Tue, 14 Jul 2026 10:12:52 +0300 Subject: [PATCH] Add retry-safe Dataproc batch ID prefix support --- .../google/docs/operators/cloud/dataproc.rst | 15 +++++ .../google/cloud/operators/dataproc.py | 54 ++++++++++++--- .../google/cloud/operators/test_dataproc.py | 67 +++++++++++++++++++ 3 files changed, 126 insertions(+), 10 deletions(-) diff --git a/providers/google/docs/operators/cloud/dataproc.rst b/providers/google/docs/operators/cloud/dataproc.rst index 115b1001eb934..6dfb771960300 100644 --- a/providers/google/docs/operators/cloud/dataproc.rst +++ b/providers/google/docs/operators/cloud/dataproc.rst @@ -506,6 +506,21 @@ Managed Spark supports creating a batch workload. A batch can be created using: :class:`~airflow.providers.google.cloud.operators.dataproc.DataprocCreateBatchOperator`. +When you need a readable ID and retry-safe uniqueness, set ``batch_id_prefix`` instead of ``batch_id``. +The operator appends a unique suffix to the prefix so retried task attempts do not fail with +``ALREADY_EXISTS`` when submitting the batch. +The prefix is used as-is (no normalization), and Dataproc validates the resulting batch ID. + +.. code-block:: python + + DataprocCreateBatchOperator( + task_id="create_batch", + project_id=PROJECT_ID, + region=REGION, + batch=BATCH_CONFIG, + batch_id_prefix="example-managed-spark-batch", + ) + The executable example below still imports the compatibility name ``DataprocCreateBatchOperator``. The preferred alias for new code is ``ManagedSparkCreateBatchOperator``. diff --git a/providers/google/src/airflow/providers/google/cloud/operators/dataproc.py b/providers/google/src/airflow/providers/google/cloud/operators/dataproc.py index e41e9b8550c5d..c97279cb4e433 100644 --- a/providers/google/src/airflow/providers/google/cloud/operators/dataproc.py +++ b/providers/google/src/airflow/providers/google/cloud/operators/dataproc.py @@ -22,6 +22,7 @@ import inspect import re import time +import uuid import warnings from collections.abc import MutableSequence, Sequence from dataclasses import dataclass @@ -2416,9 +2417,12 @@ class DataprocCreateBatchOperator(GoogleCloudBaseOperator): :param project_id: Optional. The ID of the Google Cloud project that the cluster belongs to. (templated) :param region: Required. The Cloud Dataproc region in which to handle the request. (templated) :param batch: Required. The batch to create. (templated) - :param batch_id: Required. The ID to use for the batch, which will become the final component + :param batch_id: Optional. The ID to use for the batch, which will become the final component of the batch's resource name. This value must be 4-63 characters. Valid characters are /[a-z][0-9]-/. (templated) + :param batch_id_prefix: Optional. Prefix for a generated unique batch ID. + The operator appends a random suffix to avoid collisions between retries. + Mutually exclusive with ``batch_id``. (templated) :param request_id: Optional. A unique id used to identify the request. If the server receives two ``CreateBatchRequest`` requests with the same id, then the second request will be ignored and the first ``google.longrunning.Operation`` created and stored in the backend is returned. @@ -2451,11 +2455,13 @@ class DataprocCreateBatchOperator(GoogleCloudBaseOperator): "project_id", "batch", "batch_id", + "batch_id_prefix", "region", "gcp_conn_id", "impersonation_chain", ) operator_extra_links = (DataprocBatchLink(),) + _BATCH_ID_SUFFIX_LENGTH = 8 def __init__( self, @@ -2464,6 +2470,7 @@ def __init__( project_id: str = PROVIDE_PROJECT_ID, batch: dict | Batch, batch_id: str | None = None, + batch_id_prefix: str | None = None, request_id: str | None = None, num_retries_if_resource_is_not_ready: int = 0, retry: Retry | _MethodDefault = DEFAULT, @@ -2486,10 +2493,13 @@ def __init__( super().__init__(**kwargs) if deferrable and polling_interval_seconds <= 0: raise ValueError("Invalid value for polling_interval_seconds. Expected value greater than 0") + if batch_id and batch_id_prefix: + raise ValueError("batch_id and batch_id_prefix are mutually exclusive") self.region = region self.project_id = project_id self.batch = batch self.batch_id = batch_id + self.batch_id_prefix = batch_id_prefix self.request_id = request_id self.num_retries_if_resource_is_not_ready = num_retries_if_resource_is_not_ready self.retry = retry @@ -2511,16 +2521,17 @@ def execute(self, context: Context): "Both asynchronous and deferrable parameters were passed. Please, provide only one." ) + requested_batch_id = self._resolve_requested_batch_id() batch_id: str = "" - if self.batch_id: - batch_id = self.batch_id + if requested_batch_id: + batch_id = requested_batch_id self.log.info("Starting batch %s", batch_id) # Persist the link earlier so users can observe the progress DataprocBatchLink.persist( context=context, project_id=self.project_id, region=self.region, - batch_id=self.batch_id, + batch_id=batch_id, ) else: self.log.info("Starting batch. The batch ID will be generated since it was not provided.") @@ -2536,13 +2547,15 @@ def execute(self, context: Context): region=self.region, project_id=self.project_id, batch=self.batch, - batch_id=self.batch_id, + batch_id=requested_batch_id, request_id=self.request_id, retry=self.retry, timeout=self.timeout, metadata=self.metadata, ) except AlreadyExists: + if not batch_id: + raise ValueError("Dataproc reported an existing batch without a requested batch_id.") self.log.info("Batch with given id already exists.") self.log.info("Attaching to the job %s if it is still running.", batch_id) else: @@ -2599,13 +2612,30 @@ def execute(self, context: Context): attempt = self.num_retries_if_resource_is_not_ready while attempt > 0: attempt -= 1 - batch, batch_id = self.retry_batch_creation(batch_id) + batch, batch_id = self.retry_batch_creation( + previous_batch_id=batch_id, + requested_batch_id=requested_batch_id, + ) if not self.hook.check_error_for_resource_is_not_ready_msg(batch.state_message): break self.handle_batch_status(context, batch.state.name, batch_id, batch.state_message) return Batch.to_dict(batch) + def _resolve_requested_batch_id(self) -> str | None: + if self.batch_id: + return self.batch_id + if self.batch_id_prefix: + return self._build_unique_batch_id_from_prefix(self.batch_id_prefix) + return None + + @classmethod + def _build_unique_batch_id_from_prefix(cls, batch_id_prefix: str) -> str: + prefix = str(batch_id_prefix) + if not prefix: + raise ValueError("batch_id_prefix must contain at least one valid character") + return f"{prefix}-{uuid.uuid4().hex[: cls._BATCH_ID_SUFFIX_LENGTH]}" + @cached_property def hook(self) -> DataprocHook: return DataprocHook(gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain) @@ -2647,8 +2677,9 @@ def handle_batch_status( def retry_batch_creation( self, previous_batch_id: str, + requested_batch_id: str | None, ): - self.log.info("Retrying creation process for batch_id %s", self.batch_id) + self.log.info("Retrying creation process for batch_id %s", previous_batch_id) self.log.info("Deleting previous failed Batch") self.hook.delete_batch( batch_id=previous_batch_id, @@ -2658,21 +2689,24 @@ def retry_batch_creation( timeout=self.timeout, metadata=self.metadata, ) - self.log.info("Starting a new creation for batch_id %s", self.batch_id) + self.log.info("Starting a new creation for batch_id %s", requested_batch_id) + batch_id = requested_batch_id or previous_batch_id try: self.operation = self.hook.create_batch( region=self.region, project_id=self.project_id, batch=self.batch, - batch_id=self.batch_id, + batch_id=requested_batch_id, request_id=self.request_id, retry=self.retry, timeout=self.timeout, metadata=self.metadata, ) except AlreadyExists: + if not requested_batch_id: + raise ValueError("Dataproc reported an existing batch without a requested batch_id.") self.log.info("Batch with given id already exists.") - self.log.info("Attaching to the job %s if it is still running.", self.batch_id) + self.log.info("Attaching to the job %s if it is still running.", requested_batch_id) else: if self.operation and self.operation.metadata: batch_id = self.operation.metadata.batch.split("/")[-1] diff --git a/providers/google/tests/unit/google/cloud/operators/test_dataproc.py b/providers/google/tests/unit/google/cloud/operators/test_dataproc.py index 7f0f13f77bd23..82bf8d01541e5 100644 --- a/providers/google/tests/unit/google/cloud/operators/test_dataproc.py +++ b/providers/google/tests/unit/google/cloud/operators/test_dataproc.py @@ -3737,6 +3737,73 @@ def test_missing_region_parameter(self): class TestDataprocCreateBatchOperator: + def test_batch_id_and_batch_id_prefix_are_mutually_exclusive(self): + with pytest.raises(ValueError, match="mutually exclusive"): + DataprocCreateBatchOperator( + task_id=TASK_ID, + region=GCP_REGION, + project_id=GCP_PROJECT, + batch=BATCH, + batch_id=BATCH_ID, + batch_id_prefix="batch-prefix", + ) + + @mock.patch(DATAPROC_PATH.format("uuid.uuid4")) + def test_build_unique_batch_id_from_prefix_preserves_prefix_verbatim(self, mock_uuid): + mock_uuid.return_value.hex = "abc12345def67890" + batch_id = DataprocCreateBatchOperator._build_unique_batch_id_from_prefix("Batch_Prefix") + + assert batch_id == "Batch_Prefix-abc12345" + + @mock.patch.object(DataprocCreateBatchOperator, "log", new_callable=mock.MagicMock) + @mock.patch(DATAPROC_PATH.format("uuid.uuid4")) + @mock.patch(DATAPROC_PATH.format("Batch.to_dict")) + @mock.patch(DATAPROC_PATH.format("DataprocHook")) + def test_execute_with_batch_id_prefix(self, mock_hook, to_dict_mock, mock_uuid, mock_log): + mock_uuid.return_value.hex = "abc12345def67890" + expected_batch_id = "batch-prefix-abc12345" + operator = DataprocCreateBatchOperator( + task_id=TASK_ID, + gcp_conn_id=GCP_CONN_ID, + impersonation_chain=IMPERSONATION_CHAIN, + region=GCP_REGION, + project_id=GCP_PROJECT, + batch=BATCH, + batch_id_prefix="batch-prefix", + request_id=REQUEST_ID, + retry=RETRY, + timeout=TIMEOUT, + metadata=METADATA, + ) + mock_hook.return_value.create_batch.return_value.metadata.batch = f"prefix/{expected_batch_id}" + batch_state_succeeded = Batch(state=Batch.State.SUCCEEDED) + mock_hook.return_value.wait_for_batch.return_value = batch_state_succeeded + + operator.execute(context=MagicMock()) + + mock_hook.return_value.create_batch.assert_called_once_with( + region=GCP_REGION, + project_id=GCP_PROJECT, + batch=BATCH, + batch_id=expected_batch_id, + request_id=REQUEST_ID, + retry=RETRY, + timeout=TIMEOUT, + metadata=METADATA, + ) + to_dict_mock.assert_called_once_with(batch_state_succeeded) + logs_link = DATAPROC_BATCH_LINK.format( + region=GCP_REGION, project_id=GCP_PROJECT, batch_id=expected_batch_id + ) + mock_log.info.assert_has_calls( + [ + mock.call("Starting batch %s", expected_batch_id), + mock.call("The batch %s was created.", expected_batch_id), + mock.call("Waiting for the completion of batch job %s", expected_batch_id), + mock.call("Batch job %s completed.\nDriver logs: %s", expected_batch_id, logs_link), + ] + ) + @mock.patch.object(DataprocCreateBatchOperator, "log", new_callable=mock.MagicMock) @mock.patch(DATAPROC_PATH.format("Batch.to_dict")) @mock.patch(DATAPROC_PATH.format("DataprocHook"))