You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by po...@apache.org on 2024/03/15 21:36:14 UTC
(airflow) branch main updated: Set job labels for traceability in BigQuery jobs (#37736)
This is an automated email from the ASF dual-hosted git repository.
potiuk 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 3e2bfb8b3e Set job labels for traceability in BigQuery jobs (#37736)
3e2bfb8b3e is described below
commit 3e2bfb8b3ee80ddc18b00e461de53390dcc5a8b3
Author: Pat Leamon <pa...@redbubble.com>
AuthorDate: Sat Mar 16 08:36:08 2024 +1100
Set job labels for traceability in BigQuery jobs (#37736)
* Set job labels for traceability
* Fixup
* Ensure the automatic labels will be valid
If not valid they won't be added.
* Fixup
* Merge length + regex check
* Simplify the label checking logic
---
.../providers/google/cloud/operators/bigquery.py | 17 +++
.../google/cloud/operators/test_bigquery.py | 134 +++++++++++++++++++++
2 files changed, 151 insertions(+)
diff --git a/airflow/providers/google/cloud/operators/bigquery.py b/airflow/providers/google/cloud/operators/bigquery.py
index 77862491b4..bb75b068b1 100644
--- a/airflow/providers/google/cloud/operators/bigquery.py
+++ b/airflow/providers/google/cloud/operators/bigquery.py
@@ -20,6 +20,7 @@ from __future__ import annotations
import enum
import json
+import re
import warnings
from functools import cached_property
from typing import TYPE_CHECKING, Any, Iterable, Sequence, SupportsAbs
@@ -64,6 +65,8 @@ if TYPE_CHECKING:
BIGQUERY_JOB_DETAILS_LINK_FMT = "https://console.cloud.google.com/bigquery?j={job_id}"
+LABEL_REGEX = re.compile(r"^[a-z][\w-]{0,63}$")
+
class BigQueryUIColors(enum.Enum):
"""Hex colors for BigQuery operators."""
@@ -2768,11 +2771,25 @@ class BigQueryInsertJobOperator(GoogleCloudBaseOperator, _BigQueryOpenLineageMix
with open(self.configuration) as file:
self.configuration = json.loads(file.read())
+ def _add_job_labels(self) -> None:
+ dag_label = self.dag_id.lower()
+ task_label = self.task_id.lower()
+
+ if LABEL_REGEX.match(dag_label) and LABEL_REGEX.match(task_label):
+ automatic_labels = {"airflow-dag": dag_label, "airflow-task": task_label}
+ if isinstance(self.configuration.get("labels"), dict):
+ self.configuration["labels"].update(automatic_labels)
+ elif "labels" not in self.configuration:
+ self.configuration["labels"] = automatic_labels
+
def _submit_job(
self,
hook: BigQueryHook,
job_id: str,
) -> BigQueryJob:
+ # Annotate the job with dag and task id labels
+ self._add_job_labels()
+
# Submit a new job without waiting for it to complete.
return hook.insert_job(
configuration=self.configuration,
diff --git a/tests/providers/google/cloud/operators/test_bigquery.py b/tests/providers/google/cloud/operators/test_bigquery.py
index 64e31913fb..ae0bcbfb89 100644
--- a/tests/providers/google/cloud/operators/test_bigquery.py
+++ b/tests/providers/google/cloud/operators/test_bigquery.py
@@ -1102,6 +1102,7 @@ class TestBigQueryInsertJobOperator:
project_id=TEST_GCP_PROJECT_ID,
)
result = op.execute(context=MagicMock())
+ assert configuration["labels"] == {"airflow-dag": "adhoc_airflow", "airflow-task": "insert_query_job"}
mock_hook.return_value.insert_job.assert_called_once_with(
configuration=configuration,
@@ -1143,6 +1144,7 @@ class TestBigQueryInsertJobOperator:
project_id=TEST_GCP_PROJECT_ID,
)
result = op.execute(context=MagicMock())
+ assert configuration["labels"] == {"airflow-dag": "adhoc_airflow", "airflow-task": "copy_query_job"}
mock_hook.return_value.insert_job.assert_called_once_with(
configuration=configuration,
@@ -1753,6 +1755,138 @@ class TestBigQueryInsertJobOperator:
project_id=TEST_GCP_PROJECT_ID,
)
+ @mock.patch("airflow.providers.google.cloud.operators.bigquery.BigQueryHook")
+ def test_execute_adds_to_existing_labels(self, mock_hook):
+ job_id = "123456"
+ hash_ = "hash"
+ real_job_id = f"{job_id}_{hash_}"
+
+ configuration = {
+ "query": {
+ "query": "SELECT * FROM any",
+ "useLegacySql": False,
+ },
+ "labels": {"foo": "bar"},
+ }
+ mock_hook.return_value.insert_job.return_value = MagicMock(job_id=real_job_id, error_result=False)
+ mock_hook.return_value.generate_job_id.return_value = real_job_id
+
+ op = BigQueryInsertJobOperator(
+ task_id="insert_query_job",
+ configuration=configuration,
+ location=TEST_DATASET_LOCATION,
+ job_id=job_id,
+ project_id=TEST_GCP_PROJECT_ID,
+ )
+ op.execute(context=MagicMock())
+ assert configuration["labels"] == {
+ "foo": "bar",
+ "airflow-dag": "adhoc_airflow",
+ "airflow-task": "insert_query_job",
+ }
+
+ @mock.patch("airflow.providers.google.cloud.operators.bigquery.BigQueryHook")
+ def test_execute_respects_explicit_no_labels(self, mock_hook):
+ job_id = "123456"
+ hash_ = "hash"
+ real_job_id = f"{job_id}_{hash_}"
+
+ configuration = {
+ "query": {
+ "query": "SELECT * FROM any",
+ "useLegacySql": False,
+ },
+ "labels": None,
+ }
+ mock_hook.return_value.insert_job.return_value = MagicMock(job_id=real_job_id, error_result=False)
+ mock_hook.return_value.generate_job_id.return_value = real_job_id
+
+ op = BigQueryInsertJobOperator(
+ task_id="insert_query_job",
+ configuration=configuration,
+ location=TEST_DATASET_LOCATION,
+ job_id=job_id,
+ project_id=TEST_GCP_PROJECT_ID,
+ )
+ op.execute(context=MagicMock())
+ assert configuration["labels"] is None
+
+ def test_task_label_too_big(self):
+ configuration = {
+ "query": {
+ "query": "SELECT * FROM any",
+ "useLegacySql": False,
+ },
+ }
+ op = BigQueryInsertJobOperator(
+ task_id="insert_query_job_except_this_task_id_is_really_really_really_really_long",
+ configuration=configuration,
+ location=TEST_DATASET_LOCATION,
+ project_id=TEST_GCP_PROJECT_ID,
+ )
+ op._add_job_labels()
+ assert "labels" not in configuration
+
+ def test_dag_label_too_big(self, dag_maker):
+ configuration = {
+ "query": {
+ "query": "SELECT * FROM any",
+ "useLegacySql": False,
+ },
+ }
+ with dag_maker("adhoc_airflow_except_this_task_id_is_really_really_really_really_long"):
+ op = BigQueryInsertJobOperator(
+ task_id="insert_query_job",
+ configuration=configuration,
+ location=TEST_DATASET_LOCATION,
+ project_id=TEST_GCP_PROJECT_ID,
+ )
+ op._add_job_labels()
+ assert "labels" not in configuration
+
+ def test_labels_lowercase(self, dag_maker):
+ configuration = {
+ "query": {
+ "query": "SELECT * FROM any",
+ "useLegacySql": False,
+ },
+ }
+ with dag_maker("YELLING_DAG_NAME"):
+ op = BigQueryInsertJobOperator(
+ task_id="YELLING_TASK_ID",
+ configuration=configuration,
+ location=TEST_DATASET_LOCATION,
+ project_id=TEST_GCP_PROJECT_ID,
+ )
+ op._add_job_labels()
+ assert configuration["labels"]["airflow-dag"] == "yelling_dag_name"
+ assert configuration["labels"]["airflow-task"] == "yelling_task_id"
+
+ def test_labels_invalid_names(self, dag_maker):
+ configuration = {
+ "query": {
+ "query": "SELECT * FROM any",
+ "useLegacySql": False,
+ },
+ }
+ op = BigQueryInsertJobOperator(
+ task_id="task.with.dots.is.allowed",
+ configuration=configuration,
+ location=TEST_DATASET_LOCATION,
+ project_id=TEST_GCP_PROJECT_ID,
+ )
+ op._add_job_labels()
+ assert "labels" not in configuration
+
+ op = BigQueryInsertJobOperator(
+ task_id="123_task",
+ configuration=configuration,
+ location=TEST_DATASET_LOCATION,
+ project_id=TEST_GCP_PROJECT_ID,
+ )
+ op._add_job_labels()
+ assert "labels" not in configuration
+
class TestBigQueryIntervalCheckOperator:
def test_bigquery_interval_check_operator_execute_complete(self):