This is an automated email from the ASF dual-hosted git repository.
shahar1 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 7644e35f8b7 Add retry-safe batch_id_prefix to
DataprocCreateBatchOperator (#69855)
7644e35f8b7 is described below
commit 7644e35f8b78761e4ab0c01f386f67d5e8cc0811
Author: Amir Mor <[email protected]>
AuthorDate: Sun Oct 18 09:09:14 2026 +0300
Add retry-safe batch_id_prefix to DataprocCreateBatchOperator (#69855)
Add retry-safe batch ID prefixes to DataprocCreateBatchOperator
Allow retries to create batches with distinct IDs, and document the new
option
with a runnable example.
---------
Co-authored-by: Shahar Epstein <[email protected]>
---
providers/google/docs/operators/cloud/dataproc.rst | 13 ++
.../providers/google/cloud/operators/dataproc.py | 52 +++++--
.../cloud/dataproc/example_dataproc_batch.py | 22 +++
.../unit/google/cloud/operators/test_dataproc.py | 170 +++++++++++++++++++++
4 files changed, 247 insertions(+), 10 deletions(-)
diff --git a/providers/google/docs/operators/cloud/dataproc.rst
b/providers/google/docs/operators/cloud/dataproc.rst
index 183a6b10443..056597b6399 100644
--- a/providers/google/docs/operators/cloud/dataproc.rst
+++ b/providers/google/docs/operators/cloud/dataproc.rst
@@ -512,6 +512,19 @@ The executable example below still imports the
compatibility name
:start-after: [START how_to_cloud_dataproc_create_batch_operator]
:end-before: [END how_to_cloud_dataproc_create_batch_operator]
+When you need a readable ID and retry-safe uniqueness, set ``batch_id_prefix``
instead of ``batch_id``.
+The operator generates a new ``<batch_id_prefix>-<8 hex characters>`` ID for
each task attempt, so
+retried attempts submit a new batch instead of attaching to the batch of a
previous attempt.
+``batch_id_prefix`` cannot be combined with ``request_id``, because Dataproc
returns the original
+operation for a repeated ``request_id``.
+The prefix is used as-is (no normalization), and Dataproc validates the
resulting batch ID.
+
+.. exampleinclude::
/../../google/tests/system/google/cloud/dataproc/example_dataproc_batch.py
+ :language: python
+ :dedent: 4
+ :start-after: [START
how_to_cloud_dataproc_create_batch_operator_with_prefix]
+ :end-before: [END how_to_cloud_dataproc_create_batch_operator_with_prefix]
+
For creating a batch with Persistent History Server first you should create a
Managed Spark Cluster
with specific parameters. Documentation how create cluster you can find
`here
<https://cloud.google.com/dataproc/docs/concepts/jobs/history-server#setting_up_a_persistent_history_server>`__:
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 6ef81c4bf15..9d004d8b780 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 @@ from __future__ import annotations
import inspect
import re
import time
+import uuid
from collections.abc import MutableSequence, Sequence
from dataclasses import dataclass
from datetime import datetime, timedelta
@@ -2415,9 +2416,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 batch ID generated on each
task attempt, in the form
+ ``<batch_id_prefix>-<8 hex characters>``, so a retried attempt submits
a new batch.
+ Mutually exclusive with ``batch_id`` and ``request_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.
@@ -2450,11 +2454,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,
@@ -2463,6 +2469,7 @@ class
DataprocCreateBatchOperator(GoogleCloudBaseOperator):
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,
@@ -2489,6 +2496,7 @@ class
DataprocCreateBatchOperator(GoogleCloudBaseOperator):
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
@@ -2510,16 +2518,17 @@ class
DataprocCreateBatchOperator(GoogleCloudBaseOperator):
"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.")
@@ -2535,7 +2544,7 @@ class
DataprocCreateBatchOperator(GoogleCloudBaseOperator):
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,
@@ -2598,13 +2607,33 @@ class
DataprocCreateBatchOperator(GoogleCloudBaseOperator):
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_prefix and self.request_id:
+ # Dataproc returns the original operation for a repeated
request_id, so a retried
+ # attempt would reattach to the previous batch despite its new
batch ID.
+ raise ValueError("batch_id_prefix and request_id 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:
+ return f"{batch_id_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)
@@ -2646,8 +2675,10 @@ class
DataprocCreateBatchOperator(GoogleCloudBaseOperator):
def retry_batch_creation(
self,
previous_batch_id: str,
+ requested_batch_id: str | None = None,
):
- self.log.info("Retrying creation process for batch_id %s",
self.batch_id)
+ requested_batch_id = requested_batch_id or 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,
@@ -2657,13 +2688,14 @@ class
DataprocCreateBatchOperator(GoogleCloudBaseOperator):
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,
@@ -2671,7 +2703,7 @@ class
DataprocCreateBatchOperator(GoogleCloudBaseOperator):
)
except AlreadyExists:
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/system/google/cloud/dataproc/example_dataproc_batch.py
b/providers/google/tests/system/google/cloud/dataproc/example_dataproc_batch.py
index a3fdb3c7dfa..c2c44bd136c 100644
---
a/providers/google/tests/system/google/cloud/dataproc/example_dataproc_batch.py
+++
b/providers/google/tests/system/google/cloud/dataproc/example_dataproc_batch.py
@@ -52,6 +52,7 @@ BATCH_ID = f"batch-{ENV_ID}-{DAG_ID}".replace("_", "-")
BATCH_ID_2 = f"batch-{ENV_ID}-{DAG_ID}-2".replace("_", "-")
BATCH_ID_3 = f"batch-{ENV_ID}-{DAG_ID}-3".replace("_", "-")
BATCH_ID_4 = f"batch-{ENV_ID}-{DAG_ID}-4".replace("_", "-")
+BATCH_ID_PREFIX = f"batch-{ENV_ID}-pfx".replace("_", "-")
BATCH_CONFIG = {
"spark_batch": {
@@ -101,6 +102,17 @@ with DAG(
)
# [END how_to_cloud_dataproc_create_batch_operator]
+ # [START how_to_cloud_dataproc_create_batch_operator_with_prefix]
+ create_batch_with_prefix = DataprocCreateBatchOperator(
+ task_id="create_batch_with_prefix",
+ project_id=PROJECT_ID,
+ region=REGION,
+ batch=BATCH_CONFIG,
+ batch_id_prefix=BATCH_ID_PREFIX,
+ result_retry=Retry(maximum=100.0, initial=10.0, multiplier=1.0),
+ )
+ # [END how_to_cloud_dataproc_create_batch_operator_with_prefix]
+
# [START how_to_cloud_dataproc_batch_async_sensor]
batch_async_sensor = DataprocBatchSensor(
task_id="batch_async_sensor",
@@ -158,10 +170,18 @@ with DAG(
task_id="delete_batch_4", project_id=PROJECT_ID, region=REGION,
batch_id=BATCH_ID_4
)
# [END how_to_cloud_dataproc_delete_batch_operator]
+ # The prefixed batch gets a generated suffix, so its ID is only known at
runtime.
+ delete_batch_with_prefix = DataprocDeleteBatchOperator(
+ task_id="delete_batch_with_prefix",
+ project_id=PROJECT_ID,
+ region=REGION,
+ batch_id="{{
task_instance.xcom_pull('create_batch_with_prefix')['name'].split('/') | last
}}",
+ )
delete_batch.trigger_rule = TriggerRule.ALL_DONE
delete_batch_2.trigger_rule = TriggerRule.ALL_DONE
delete_batch_3.trigger_rule = TriggerRule.ALL_DONE
delete_batch_4.trigger_rule = TriggerRule.ALL_FAILED
+ delete_batch_with_prefix.trigger_rule = TriggerRule.ALL_DONE
(
# TEST SETUP
@@ -172,6 +192,7 @@ with DAG(
>> batch_async_sensor
>> get_batch
>> list_batches
+ >> create_batch_with_prefix
>> create_batch_4
>> cancel_operation
# TEST TEARDOWN
@@ -179,6 +200,7 @@ with DAG(
>> delete_batch_2
>> delete_batch_3
>> delete_batch_4
+ >> delete_batch_with_prefix
)
from tests_common.test_utils.watcher import watcher
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 c1cb55f566f..c8b61bfb12f 100644
--- a/providers/google/tests/unit/google/cloud/operators/test_dataproc.py
+++ b/providers/google/tests/unit/google/cloud/operators/test_dataproc.py
@@ -3742,6 +3742,176 @@ class TestDataprocCreateWorkflowTemplateOperator:
class TestDataprocCreateBatchOperator:
+ @pytest.mark.parametrize(
+ ("extra_kwargs", "match"),
+ [
+ pytest.param({"batch_id": BATCH_ID}, "batch_id and
batch_id_prefix", id="batch_id"),
+ pytest.param({"request_id": REQUEST_ID}, "batch_id_prefix and
request_id", id="request_id"),
+ ],
+ )
+ @mock.patch(DATAPROC_PATH.format("DataprocHook"))
+ def test_batch_id_prefix_mutually_exclusive_params(self, mock_hook,
extra_kwargs, match):
+ operator = DataprocCreateBatchOperator(
+ task_id=TASK_ID,
+ region=GCP_REGION,
+ project_id=GCP_PROJECT,
+ batch=BATCH,
+ batch_id_prefix="batch-prefix",
+ **extra_kwargs,
+ )
+ with pytest.raises(ValueError, match=match):
+ 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",
+ 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=None,
+ 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(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_generates_new_id_per_task_attempt(
+ self, mock_hook, to_dict_mock, mock_uuid
+ ):
+ mock_uuid.side_effect = [mock.Mock(hex="aaaaaaaa1111"),
mock.Mock(hex="bbbbbbbb2222")]
+ operator = DataprocCreateBatchOperator(
+ task_id=TASK_ID,
+ region=GCP_REGION,
+ project_id=GCP_PROJECT,
+ batch=BATCH,
+ batch_id_prefix="batch-prefix",
+ )
+ mock_hook.return_value.wait_for_batch.return_value =
Batch(state=Batch.State.SUCCEEDED)
+
+ operator.execute(context=MagicMock())
+ operator.execute(context=MagicMock())
+
+ assert [c.kwargs["batch_id"] for c in
mock_hook.return_value.create_batch.call_args_list] == [
+ "batch-prefix-aaaaaaaa",
+ "batch-prefix-bbbbbbbb",
+ ]
+
+ @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_retries_resource_not_ready_with_same_id(
+ self, mock_hook, to_dict_mock, mock_uuid
+ ):
+ mock_uuid.return_value.hex = "abc12345def67890"
+ expected_batch_id = "batch-prefix-abc12345"
+ operator = DataprocCreateBatchOperator(
+ task_id=TASK_ID,
+ region=GCP_REGION,
+ project_id=GCP_PROJECT,
+ batch=BATCH,
+ batch_id_prefix="batch-prefix",
+ num_retries_if_resource_is_not_ready=1,
+ retry=RETRY,
+ timeout=TIMEOUT,
+ metadata=METADATA,
+ )
+ mock_hook.return_value.create_batch.return_value.metadata.batch =
f"prefix/{expected_batch_id}"
+ mock_hook.return_value.wait_for_batch.side_effect = [
+ Batch(state=Batch.State.FAILED, state_message="resource not
ready"),
+ Batch(state=Batch.State.SUCCEEDED),
+ ]
+
mock_hook.return_value.check_error_for_resource_is_not_ready_msg.side_effect =
[True, False]
+
+ operator.execute(context=MagicMock())
+
+ mock_uuid.assert_called_once()
+ mock_hook.return_value.delete_batch.assert_called_once_with(
+ batch_id=expected_batch_id,
+ region=GCP_REGION,
+ project_id=GCP_PROJECT,
+ retry=RETRY,
+ timeout=TIMEOUT,
+ metadata=METADATA,
+ )
+ assert [c.kwargs["batch_id"] for c in
mock_hook.return_value.create_batch.call_args_list] == [
+ expected_batch_id,
+ expected_batch_id,
+ ]
+
+ @mock.patch(DATAPROC_PATH.format("DataprocHook"))
+ def test_retry_batch_creation_already_exists_attaches_to_batch_id(self,
mock_hook):
+ op = DataprocCreateBatchOperator(
+ task_id=TASK_ID,
+ region=GCP_REGION,
+ project_id=GCP_PROJECT,
+ batch=BATCH,
+ batch_id=BATCH_ID,
+ retry=RETRY,
+ timeout=TIMEOUT,
+ metadata=METADATA,
+ )
+ mock_hook.return_value.create_batch.side_effect = AlreadyExists("")
+ mock_hook.return_value.wait_for_batch.return_value =
Batch(state=Batch.State.SUCCEEDED)
+
+ _, batch_id = op.retry_batch_creation(BATCH_ID)
+
+ assert batch_id == BATCH_ID
+ assert
mock_hook.return_value.create_batch.call_args.kwargs["batch_id"] == BATCH_ID
+ mock_hook.return_value.wait_for_batch.assert_called_once_with(
+ batch_id=BATCH_ID,
+ region=GCP_REGION,
+ project_id=GCP_PROJECT,
+ retry=RETRY,
+ timeout=TIMEOUT,
+ metadata=METADATA,
+ )
+
@mock.patch.object(DataprocCreateBatchOperator, "log",
new_callable=mock.MagicMock)
@mock.patch(DATAPROC_PATH.format("Batch.to_dict"))
@mock.patch(DATAPROC_PATH.format("DataprocHook"))
