Skip to content

Set task/dag labels for Dataproc Batch #46781

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
eladkal merged 2 commits into from
Mar 9, 2025
Merged

Set task/dag labels for Dataproc Batch #46781

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
630f77b
Set task/dag labels for Dataproc Batch
ntr Feb 14, 2025
3a25cd3
Adding airflow-dag-display-name label support
ntr Feb 26, 2025
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
27 changes: 27 additions & 0 deletions providers/google/src/airflow/providers/google/cloud/operators/dataproc.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2530,6 +2530,8 @@ def execute(self, context: Context):
self.log.info("Automatic injection of OpenLineage information into Spark properties is enabled.")
self._inject_openlineage_properties_into_dataproc_batch(context)

self.__update_batch_labels()

try:
self.operation = self.hook.create_batch(
region=self.region,
Expand Down Expand Up @@ -2710,6 +2712,31 @@ def _inject_openlineage_properties_into_dataproc_batch(self, context: Context) -
exc_info=e,
)

def __update_batch_labels(self):
dag_id = re.sub(r"[.\s]", "_", self.dag_id.lower())
task_id = re.sub(r"[.\s]", "_", self.task_id.lower())

labels_regex = re.compile(r"^[a-z][\w-]{0,63}$")
if not labels_regex.match(dag_id) or not labels_regex.match(task_id):
return

labels_limit = 32
new_labels = {"airflow-dag-id": dag_id, "airflow-task-id": task_id}

if self._dag:
dag_display_name = re.sub(r"[.\s]", "_", self._dag.dag_display_name.lower())
if labels_regex.match(dag_id):
new_labels["airflow-dag-display-name"] = dag_display_name

if isinstance(self.batch, Batch):
if len(self.batch.labels) + len(new_labels) <= labels_limit:
self.batch.labels.update(new_labels)
elif "labels" not in self.batch:
self.batch["labels"] = new_labels
elif isinstance(self.batch.get("labels"), dict):
if len(self.batch["labels"]) + len(new_labels) <= labels_limit:
self.batch["labels"].update(new_labels)


class DataprocDeleteBatchOperator(GoogleCloudBaseOperator):
"""
Expand Down
108 changes: 107 additions & 1 deletion providers/google/tests/unit/google/cloud/operators/test_dataproc.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,8 +18,9 @@

import datetime as dt
import inspect
from copy import deepcopy
from unittest import mock
from unittest.mock import MagicMock, Mock, call
from unittest.mock import ANY, MagicMock, Mock, call

import pytest
from google.api_core.exceptions import AlreadyExists, NotFound
Expand Down Expand Up @@ -3775,6 +3776,111 @@ def test_execute_openlineage_transport_info_injection_skipped_when_ol_not_access
metadata=METADATA,
)

@staticmethod
def __assert_batch_create(mock_hook, expected_batch):
mock_hook.return_value.create_batch.assert_called_once_with(
region=ANY,
project_id=ANY,
batch=expected_batch,
batch_id=ANY,
request_id=ANY,
retry=ANY,
timeout=ANY,
metadata=ANY,
)

@mock.patch(DATAPROC_PATH.format("Batch.to_dict"))
@mock.patch(DATAPROC_PATH.format("DataprocHook"))
def test_create_batch_asdict_labels_updated(self, mock_hook, to_dict_mock):
expected_labels = {
"airflow-dag-id": "test_dag",
"airflow-dag-display-name": "test_dag",
"airflow-task-id": "test-task",
}

expected_batch = {
**BATCH,
"labels": expected_labels,
}

DataprocCreateBatchOperator(
task_id="test-task",
dag=DAG(dag_id="test_dag"),
batch=BATCH,
region=GCP_REGION,
).execute(context=EXAMPLE_CONTEXT)

TestDataprocCreateBatchOperator.__assert_batch_create(mock_hook, expected_batch)

@mock.patch(DATAPROC_PATH.format("Batch.to_dict"))
@mock.patch(DATAPROC_PATH.format("DataprocHook"))
def test_create_batch_asdict_labels_uppercase_transformed(self, mock_hook, to_dict_mock):
expected_labels = {
"airflow-dag-id": "test_dag",
"airflow-dag-display-name": "test_dag",
"airflow-task-id": "test-task",
}

expected_batch = {
**BATCH,
"labels": expected_labels,
}

DataprocCreateBatchOperator(
task_id="test-TASK",
dag=DAG(dag_id="Test_dag"),
batch=BATCH,
region=GCP_REGION,
).execute(context=EXAMPLE_CONTEXT)

TestDataprocCreateBatchOperator.__assert_batch_create(mock_hook, expected_batch)

@mock.patch(DATAPROC_PATH.format("Batch.to_dict"))
@mock.patch(DATAPROC_PATH.format("DataprocHook"))
def test_create_batch_invalid_taskid_labels_ignored(self, mock_hook, to_dict_mock):
DataprocCreateBatchOperator(
task_id=".task-id",
dag=DAG(dag_id="test-dag"),
batch=BATCH,
region=GCP_REGION,
).execute(context=EXAMPLE_CONTEXT)

TestDataprocCreateBatchOperator.__assert_batch_create(mock_hook, BATCH)

@mock.patch(DATAPROC_PATH.format("Batch.to_dict"))
@mock.patch(DATAPROC_PATH.format("DataprocHook"))
def test_create_batch_long_taskid_labels_ignored(self, mock_hook, to_dict_mock):
DataprocCreateBatchOperator(
task_id="a" * 65,
dag=DAG(dag_id="test-dag"),
batch=BATCH,
region=GCP_REGION,
).execute(context=EXAMPLE_CONTEXT)

TestDataprocCreateBatchOperator.__assert_batch_create(mock_hook, BATCH)

@mock.patch(DATAPROC_PATH.format("Batch.to_dict"))
@mock.patch(DATAPROC_PATH.format("DataprocHook"))
def test_create_batch_asobj_labels_updated(self, mock_hook, to_dict_mock):
batch = Batch(name="test")
batch.labels["foo"] = "bar"
dag = DAG(dag_id="test_dag")

expected_labels = {
"airflow-dag-id": "test_dag",
"airflow-dag-display-name": "test_dag",
"airflow-task-id": "test-task",
}

expected_batch = deepcopy(batch)
expected_batch.labels.update(expected_labels)

DataprocCreateBatchOperator(task_id="test-task", batch=batch, region=GCP_REGION, dag=dag).execute(
context=EXAMPLE_CONTEXT
)

TestDataprocCreateBatchOperator.__assert_batch_create(mock_hook, expected_batch)


class TestDataprocDeleteBatchOperator:
@mock.patch(DATAPROC_PATH.format("DataprocHook"))
Expand Down