Google: Add retry-safe batch_id_prefix to DataprocCreateBatchOperator

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``.

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,
@@ -2490,6 +2497,7 @@ def __init__(
         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 +2519,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 +2545,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 +2610,32 @@ 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 and self.batch_id_prefix:
+            raise ValueError("batch_id and batch_id_prefix are mutually exclusive")
+        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]

providers/google/tests/unit/google/cloud/operators/test_dataproc.py

@@ -3737,6 +3737,76 @@ def test_missing_region_parameter(self):
 
 
 class TestDataprocCreateBatchOperator:
+    @mock.patch(DATAPROC_PATH.format("DataprocHook"))
+    def test_batch_id_and_batch_id_prefix_are_mutually_exclusive(self, mock_hook):
+        operator = DataprocCreateBatchOperator(
+            task_id=TASK_ID,
+            region=GCP_REGION,
+            project_id=GCP_PROJECT,
+            batch=BATCH,
+            batch_id=BATCH_ID,
+            batch_id_prefix="batch-prefix",
+        )
+        with pytest.raises(ValueError, match="mutually exclusive"):
+            operator.execute(context=MagicMock())
+        mock_hook.return_value.create_batch.assert_not_called()
+
+    @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"))