apache-airflow-providers-google 10.16.0rc1__py3-none-any.whl → 10.17.0__py3-none-any.whl

airflow/providers/google/cloud/operators/bigquery.py

@@ -16,10 +16,12 @@
 # specific language governing permissions and limitations
 # under the License.
 """This module contains Google BigQuery operators."""
+
 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 +66,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."""
@@ -1213,6 +1217,7 @@ class BigQueryExecuteQueryOperator(GoogleCloudBaseOperator):
         location: str | None = None,
         encryption_configuration: dict | None = None,
         impersonation_chain: str | Sequence[str] | None = None,
+        impersonation_scopes: str | Sequence[str] | None = None,
         job_id: str | list[str] | None = None,
         **kwargs,
     ) -> None:
@@ -1239,6 +1244,7 @@ class BigQueryExecuteQueryOperator(GoogleCloudBaseOperator):
         self.encryption_configuration = encryption_configuration
         self.hook: BigQueryHook | None = None
         self.impersonation_chain = impersonation_chain
+        self.impersonation_scopes = impersonation_scopes
         self.job_id = job_id
 
     def execute(self, context: Context):
@@ -1249,6 +1255,7 @@ class BigQueryExecuteQueryOperator(GoogleCloudBaseOperator):
                 use_legacy_sql=self.use_legacy_sql,
                 location=self.location,
                 impersonation_chain=self.impersonation_chain,
+                impersonation_scopes=self.impersonation_scopes,
             )
         if isinstance(self.sql, str):
             self.job_id = self.hook.run_query(
@@ -2768,11 +2775,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,

airflow/providers/google/cloud/operators/bigquery_dts.py

@@ -16,6 +16,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """This module contains Google BigQuery Data Transfer Service operators."""
+
 from __future__ import annotations
 
 import time

airflow/providers/google/cloud/operators/bigtable.py

@@ -16,6 +16,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """This module contains Google Cloud Bigtable operators."""
+
 from __future__ import annotations
 
 from typing import TYPE_CHECKING, Iterable, Sequence

airflow/providers/google/cloud/operators/cloud_base.py

@@ -16,6 +16,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """This module contains a Google API base operator."""
+
 from __future__ import annotations
 
 from google.api_core.gapic_v1.method import DEFAULT

airflow/providers/google/cloud/operators/cloud_build.py

@@ -16,6 +16,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """Operators that integrates with Google Cloud Build service."""
+
 from __future__ import annotations
 
 import json

airflow/providers/google/cloud/operators/cloud_memorystore.py

@@ -23,6 +23,7 @@ Operators for Google Cloud Memorystore service.
     FieldMask
     memcache
 """
+
 from __future__ import annotations
 
 from typing import TYPE_CHECKING, Sequence

airflow/providers/google/cloud/operators/cloud_sql.py

@@ -16,6 +16,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """This module contains Google Cloud SQL operators."""
+
 from __future__ import annotations
 
 from typing import TYPE_CHECKING, Any, Iterable, Mapping, Sequence

airflow/providers/google/cloud/operators/cloud_storage_transfer_service.py

@@ -16,6 +16,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """This module contains Google Cloud Transfer operators."""
+
 from __future__ import annotations
 
 from copy import deepcopy
@@ -530,7 +531,7 @@ class CloudDataTransferServiceListOperationsOperator(GoogleCloudBaseOperator):
 
     # [START gcp_transfer_operations_list_template_fields]
     template_fields: Sequence[str] = (
-        "filter",
+        "request_filter",
         "gcp_conn_id",
         "google_impersonation_chain",
     )
@@ -547,22 +548,27 @@ class CloudDataTransferServiceListOperationsOperator(GoogleCloudBaseOperator):
         **kwargs,
     ) -> None:
         super().__init__(**kwargs)
-        self.filter = request_filter
+        self.request_filter = request_filter
         self.project_id = project_id
         self.gcp_conn_id = gcp_conn_id
         self.api_version = api_version
         self.google_impersonation_chain = google_impersonation_chain
 
+    @property
+    def filter(self) -> dict | None:
+        """Alias for ``request_filter``, used for compatibility."""
+        return self.request_filter
+
     def execute(self, context: Context) -> list[dict]:
-        if not self.filter:
-            raise AirflowException("The required parameter 'filter' is empty or None")
+        if not self.request_filter:
+            raise AirflowException("The required parameter 'request_filter' is empty or None")
 
         hook = CloudDataTransferServiceHook(
             api_version=self.api_version,
             gcp_conn_id=self.gcp_conn_id,
             impersonation_chain=self.google_impersonation_chain,
         )
-        operations_list = hook.list_transfer_operations(request_filter=self.filter)
+        operations_list = hook.list_transfer_operations(request_filter=self.request_filter)
         self.log.info(operations_list)
 
         project_id = self.project_id or hook.project_id

airflow/providers/google/cloud/operators/compute.py

@@ -16,6 +16,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """This module contains Google Compute Engine operators."""
+
 from __future__ import annotations
 
 from copy import deepcopy

airflow/providers/google/cloud/operators/dataflow.py

@@ -16,6 +16,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """This module contains Google Dataflow operators."""
+
 from __future__ import annotations
 
 import copy

airflow/providers/google/cloud/operators/datafusion.py

@@ -15,6 +15,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """This module contains Google DataFusion operators."""
+
 from __future__ import annotations
 
 import time

airflow/providers/google/cloud/operators/datapipeline.py

@@ -16,6 +16,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """This module contains Google Data Pipelines operators."""
+
 from __future__ import annotations
 
 from typing import TYPE_CHECKING, Sequence

airflow/providers/google/cloud/operators/dataprep.py

@@ -16,6 +16,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """This module contains a Google Dataprep operator."""
+
 from __future__ import annotations
 
 from typing import TYPE_CHECKING, Sequence

airflow/providers/google/cloud/operators/dataproc.py

@@ -16,6 +16,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """This module contains Google Dataproc operators."""
+
 from __future__ import annotations
 
 import inspect
@@ -1429,13 +1430,13 @@ class DataprocJobBaseOperator(GoogleCloudBaseOperator):
         if self.job_template:
             job = self.job_template.build()
             return job["job"]
-        raise Exception("Create a job template before")
+        raise AirflowException("Create a job template before")
 
     def execute(self, context: Context):
         if self.job_template:
             self.job = self.job_template.build()
             if self.job is None:
-                raise Exception("The job should be set here.")
+                raise AirflowException("The job should be set here.")
             self.dataproc_job_id = self.job["job"]["reference"]["job_id"]
             self.log.info("Submitting %s job %s", self.job_type, self.dataproc_job_id)
             job_object = self.hook.submit_job(

airflow/providers/google/cloud/operators/dataproc_metastore.py

@@ -16,6 +16,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """This module contains Google Dataproc Metastore operators."""
+
 from __future__ import annotations
 
 import time

airflow/providers/google/cloud/operators/datastore.py

@@ -16,6 +16,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """This module contains Google Datastore operators."""
+
 from __future__ import annotations
 
 from typing import TYPE_CHECKING, Any, Sequence

airflow/providers/google/cloud/operators/functions.py

@@ -16,6 +16,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """This module contains Google Cloud Functions operators."""
+
 from __future__ import annotations
 
 import re

airflow/providers/google/cloud/operators/gcs.py

@@ -16,6 +16,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """This module contains a Google Cloud Storage Bucket operator."""
+
 from __future__ import annotations
 
 import datetime