From ab5dcef8c4138b34cd8ff8ea5778180f74882d78 Mon Sep 17 00:00:00 2001
From: Issam <kalliathissam@gmail.com>
Date: Sat, 11 May 2024 18:06:06 +0000
Subject: [PATCH] Fix validation for BQ label values for
 BigQueryInsertJobOperator

---
 .../google/cloud/operators/bigquery.py        |  2 +-
 .../google/cloud/operators/test_bigquery.py   | 58 ++++++++++++++++++-
 2 files changed, 58 insertions(+), 2 deletions(-)

diff --git a/airflow/providers/google/cloud/operators/bigquery.py b/airflow/providers/google/cloud/operators/bigquery.py
index dcd971ad0b221..9be3a0c2eef18 100644
--- a/airflow/providers/google/cloud/operators/bigquery.py
+++ b/airflow/providers/google/cloud/operators/bigquery.py
@@ -68,7 +68,7 @@
 
 BIGQUERY_JOB_DETAILS_LINK_FMT = "https://console.cloud.google.com/bigquery?j={job_id}"
 
-LABEL_REGEX = re.compile(r"^[a-z][\w-]{0,63}$")
+LABEL_REGEX = re.compile(r"^[\w-]{0,63}$")
 
 
 class BigQueryUIColors(enum.Enum):
diff --git a/tests/providers/google/cloud/operators/test_bigquery.py b/tests/providers/google/cloud/operators/test_bigquery.py
index ba94347437eef..d84218f8bd5e0 100644
--- a/tests/providers/google/cloud/operators/test_bigquery.py
+++ b/tests/providers/google/cloud/operators/test_bigquery.py
@@ -1921,6 +1921,62 @@ def test_labels_lowercase(self, dag_maker):
         assert configuration["labels"]["airflow-dag"] == "yelling_dag_name"
         assert configuration["labels"]["airflow-task"] == "yelling_task_id"
 
+    def test_labels_starting_with_numbers(self, dag_maker):
+        configuration = {
+            "query": {
+                "query": "SELECT * FROM any",
+                "useLegacySql": False,
+            },
+        }
+        with dag_maker("123_dag"):
+            op = BigQueryInsertJobOperator(
+                task_id="123_task",
+                configuration=configuration,
+                location=TEST_DATASET_LOCATION,
+                project_id=TEST_GCP_PROJECT_ID,
+            )
+        op._add_job_labels()
+        assert configuration["labels"]["airflow-dag"] == "123_dag"
+        assert configuration["labels"]["airflow-task"] == "123_task"
+
+    def test_labels_starting_with_underscore(self, dag_maker):
+        configuration = {
+            "query": {
+                "query": "SELECT * FROM any",
+                "useLegacySql": False,
+            },
+        }
+        with dag_maker("_dag_starting_with_underscore"):
+            op = BigQueryInsertJobOperator(
+                task_id="_task_starting_with_underscore",
+                configuration=configuration,
+                location=TEST_DATASET_LOCATION,
+                project_id=TEST_GCP_PROJECT_ID,
+            )
+        op._add_job_labels()
+        assert "labels" in configuration
+        assert configuration["labels"]["airflow-dag"] == "_dag_starting_with_underscore"
+        assert configuration["labels"]["airflow-task"] == "_task_starting_with_underscore"
+
+    def test_labels_starting_with_hyphen(self, dag_maker):
+        configuration = {
+            "query": {
+                "query": "SELECT * FROM any",
+                "useLegacySql": False,
+            },
+        }
+        with dag_maker("-dag-starting-with-hyphen"):
+            op = BigQueryInsertJobOperator(
+                task_id="-task-starting-with-hyphen",
+                configuration=configuration,
+                location=TEST_DATASET_LOCATION,
+                project_id=TEST_GCP_PROJECT_ID,
+            )
+        op._add_job_labels()
+        assert "labels" in configuration
+        assert configuration["labels"]["airflow-dag"] == "-dag-starting-with-hyphen"
+        assert configuration["labels"]["airflow-task"] == "-task-starting-with-hyphen"
+
     def test_labels_invalid_names(self, dag_maker):
         configuration = {
             "query": {
@@ -1938,7 +1994,7 @@ def test_labels_invalid_names(self, dag_maker):
         assert "labels" not in configuration
 
         op = BigQueryInsertJobOperator(
-            task_id="123_task",
+            task_id="task_id_with_exactly_64_characters_00000000000000000000000000000",
             configuration=configuration,
             location=TEST_DATASET_LOCATION,
             project_id=TEST_GCP_PROJECT_ID,