This is an automated email from the ASF dual-hosted git repository.
kaxil 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 0bcec1d359d Fix passing labels to BQ job (#65663)
0bcec1d359d is described below
commit 0bcec1d359da15dd092cc8da4bfc33bb8cf10948
Author: Ulada Zakharava <[email protected]>
AuthorDate: Wed Apr 22 12:39:16 2026 +0000
Fix passing labels to BQ job (#65663)
---
.../providers/google/cloud/operators/bigquery.py | 9 +++-
.../unit/google/cloud/operators/test_bigquery.py | 52 ++++++++++++++++++++++
2 files changed, 60 insertions(+), 1 deletion(-)
diff --git
a/providers/google/src/airflow/providers/google/cloud/operators/bigquery.py
b/providers/google/src/airflow/providers/google/cloud/operators/bigquery.py
index 85c32a540dc..c00dc2cebc5 100644
--- a/providers/google/src/airflow/providers/google/cloud/operators/bigquery.py
+++ b/providers/google/src/airflow/providers/google/cloud/operators/bigquery.py
@@ -2340,7 +2340,12 @@ class BigQueryInsertJobOperator(GoogleCloudBaseOperator,
_BigQueryInsertJobOpera
with open(self.configuration) as file:
self.configuration = json.loads(file.read())
- def _add_job_labels(self) -> None:
+ def _add_job_labels(self, hook: BigQueryHook | None = None) -> None:
+ if hook and hook.labels and "labels" not in self.configuration:
+ self.configuration["labels"] = dict(hook.labels)
+ elif hook and hook.labels and
isinstance(self.configuration.get("labels"), dict):
+ self.configuration["labels"] = {**hook.labels,
**self.configuration["labels"]}
+
dag_label = self.dag_id.lower()
task_label = self.task_id.lower().replace(".", "-")
@@ -2409,6 +2414,8 @@ class BigQueryInsertJobOperator(GoogleCloudBaseOperator,
_BigQueryInsertJobOpera
if self.project_id is None:
self.project_id = hook.project_id
+ self._add_job_labels(hook)
+
# Handles Operator retries when a user does not explicitly set a
job_id.
# For example, if a previous job failed due to a 429 "Too Many
Requests" error,
# the Operator will retry and resubmit the job. We need to ensure we
don't lose
diff --git
a/providers/google/tests/unit/google/cloud/operators/test_bigquery.py
b/providers/google/tests/unit/google/cloud/operators/test_bigquery.py
index 2655132009b..70427e794e6 100644
--- a/providers/google/tests/unit/google/cloud/operators/test_bigquery.py
+++ b/providers/google/tests/unit/google/cloud/operators/test_bigquery.py
@@ -1890,6 +1890,7 @@ class TestBigQueryInsertJobOperator:
},
"labels": None,
}
+ mock_hook.return_value.labels = {"airflow-env": "test-env-2"}
mock_hook.return_value.insert_job.return_value = MagicMock(
state="DONE", job_id=real_job_id, error_result=False
)
@@ -1905,6 +1906,57 @@ class TestBigQueryInsertJobOperator:
op.execute(context=MagicMock())
assert configuration["labels"] is None
+
@mock.patch("airflow.providers.google.cloud.operators.bigquery.BigQueryHook")
+ def test_execute_inherits_labels_from_connection(self, mock_hook):
+ job_id = "123456"
+ hash_ = "hash"
+ real_job_id = f"{job_id}_{hash_}"
+
+ configuration = {
+ "query": {
+ "query": "SELECT * FROM any",
+ "useLegacySql": False,
+ },
+ }
+ mock_hook.return_value.labels = {"airflow-env": "test-env-2"}
+ mock_hook.return_value.insert_job.return_value = MagicMock(
+ state="DONE", 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"]["airflow-env"] == "test-env-2"
+ assert configuration["labels"]["airflow-task"] == "insert_query_job"
+
+ def test_add_job_labels_merges_connection_labels(self):
+ configuration = {
+ "query": {
+ "query": "SELECT * FROM any",
+ "useLegacySql": False,
+ },
+ "labels": {"manual-label": "manual-value"},
+ }
+ op = BigQueryInsertJobOperator(
+ task_id="insert_query_job",
+ configuration=configuration,
+ location=TEST_DATASET_LOCATION,
+ project_id=TEST_GCP_PROJECT_ID,
+ )
+
+ op._add_job_labels(hook=MagicMock(labels={"airflow-env":
"test-env-2"}))
+
+ assert configuration["labels"]["airflow-env"] == "test-env-2"
+ assert configuration["labels"]["manual-label"] == "manual-value"
+ assert configuration["labels"]["airflow-task"] == "insert_query_job"
+
def test_task_label_too_big(self):
configuration = {
"query": {
