apache-airflow-providers-google 14.0.0__py3-none-any.whl → 19.1.0rc1__py3-none-any.whl

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

@@ -30,6 +30,11 @@ from enum import Enum
 from functools import cached_property
 from typing import TYPE_CHECKING, Any
 
+from google.api_core.exceptions import AlreadyExists, NotFound
+from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
+from google.api_core.retry import Retry, exponential_sleep_generator
+from google.cloud.dataproc_v1 import Batch, Cluster, ClusterStatus, JobStatus
+
 from airflow.configuration import conf
 from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
 from airflow.providers.google.cloud.hooks.dataproc import (
@@ -39,7 +44,6 @@ from airflow.providers.google.cloud.hooks.dataproc import (
 )
 from airflow.providers.google.cloud.links.dataproc import (
     DATAPROC_BATCH_LINK,
-    DATAPROC_CLUSTER_LINK_DEPRECATED,
     DATAPROC_JOB_LINK_DEPRECATED,
     DataprocBatchesListLink,
     DataprocBatchLink,
@@ -58,22 +62,18 @@ from airflow.providers.google.cloud.triggers.dataproc import (
     DataprocSubmitTrigger,
 )
 from airflow.providers.google.cloud.utils.dataproc import DataprocOperationType
-from airflow.providers.google.common.deprecated import deprecated
 from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
 from airflow.utils import timezone
-from google.api_core.exceptions import AlreadyExists, NotFound
-from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
-from google.api_core.retry import Retry, exponential_sleep_generator
-from google.cloud.dataproc_v1 import Batch, Cluster, ClusterStatus, JobStatus
 
 if TYPE_CHECKING:
-    from airflow.utils.context import Context
     from google.api_core import operation
     from google.api_core.retry_async import AsyncRetry
     from google.protobuf.duration_pb2 import Duration
     from google.protobuf.field_mask_pb2 import FieldMask
     from google.type.interval_pb2 import Interval
 
+    from airflow.providers.common.compat.sdk import Context
+
 
 class PreemptibilityType(Enum):
     """Contains possible Type values of Preemptibility applicable for every secondary worker of Cluster."""
@@ -213,6 +213,7 @@ class ClusterGenerator:
     :param secondary_worker_accelerator_type: Type of the accelerator card (GPU) to attach to the secondary workers,
         see https://cloud.google.com/dataproc/docs/reference/rest/v1/InstanceGroupConfig#acceleratorconfig
     :param secondary_worker_accelerator_count: Number of accelerator cards (GPUs) to attach to the secondary workers
+    :param cluster_tier: The tier of the cluster (e.g. "CLUSTER_TIER_STANDARD" / "CLUSTER_TIER_PREMIUM").
     """
 
     def __init__(
@@ -261,6 +262,8 @@ class ClusterGenerator:
         secondary_worker_instance_flexibility_policy: InstanceFlexibilityPolicy | None = None,
         secondary_worker_accelerator_type: str | None = None,
         secondary_worker_accelerator_count: int | None = None,
+        *,
+        cluster_tier: str | None = None,
         **kwargs,
     ) -> None:
         self.project_id = project_id
@@ -308,6 +311,7 @@ class ClusterGenerator:
         self.secondary_worker_instance_flexibility_policy = secondary_worker_instance_flexibility_policy
         self.secondary_worker_accelerator_type = secondary_worker_accelerator_type
         self.secondary_worker_accelerator_count = secondary_worker_accelerator_count
+        self.cluster_tier = cluster_tier
 
         if self.custom_image and self.image_version:
             raise ValueError("The custom_image and image_version can't be both set")
@@ -340,7 +344,7 @@ class ClusterGenerator:
             unit = match.group(2)
             if unit == "s":
                 return {"seconds": val}
-            elif unit == "m":
+            if unit == "m":
                 return {"seconds": int(timedelta(minutes=val).total_seconds())}
 
         raise AirflowException(
@@ -513,6 +517,9 @@ class ClusterGenerator:
         if self.driver_pool_size > 0:
             cluster_data["auxiliary_node_groups"] = [self._build_driver_pool()]
 
+        if self.cluster_tier:
+            cluster_data["cluster_tier"] = self.cluster_tier
+
         cluster_data = self._build_gce_cluster_config(cluster_data)
 
         if self.single_node:
@@ -621,6 +628,7 @@ class DataprocCreateClusterOperator(GoogleCloudBaseOperator):
         "virtual_cluster_config",
         "cluster_name",
         "labels",
+        "gcp_conn_id",
         "impersonation_chain",
     )
     template_fields_renderers = {"cluster_config": "json", "virtual_cluster_config": "json"}
@@ -807,7 +815,6 @@ class DataprocCreateClusterOperator(GoogleCloudBaseOperator):
         if project_id:
             DataprocClusterLink.persist(
                 context=context,
-                operator=self,
                 cluster_id=self.cluster_name,
                 project_id=project_id,
                 region=self.region,
@@ -822,26 +829,24 @@ class DataprocCreateClusterOperator(GoogleCloudBaseOperator):
                 )
                 self.log.info("Cluster created.")
                 return Cluster.to_dict(cluster)
-            else:
-                cluster = hook.get_cluster(
-                    project_id=self.project_id, region=self.region, cluster_name=self.cluster_name
-                )
-                if cluster.status.state == cluster.status.State.RUNNING:
-                    self.log.info("Cluster created.")
-                    return Cluster.to_dict(cluster)
-                else:
-                    self.defer(
-                        trigger=DataprocClusterTrigger(
-                            cluster_name=self.cluster_name,
-                            project_id=self.project_id,
-                            region=self.region,
-                            gcp_conn_id=self.gcp_conn_id,
-                            impersonation_chain=self.impersonation_chain,
-                            polling_interval_seconds=self.polling_interval_seconds,
-                            delete_on_error=self.delete_on_error,
-                        ),
-                        method_name="execute_complete",
-                    )
+            cluster = hook.get_cluster(
+                project_id=self.project_id, region=self.region, cluster_name=self.cluster_name
+            )
+            if cluster.status.state == cluster.status.State.RUNNING:
+                self.log.info("Cluster created.")
+                return Cluster.to_dict(cluster)
+            self.defer(
+                trigger=DataprocClusterTrigger(
+                    cluster_name=self.cluster_name,
+                    project_id=self.project_id,
+                    region=self.region,
+                    gcp_conn_id=self.gcp_conn_id,
+                    impersonation_chain=self.impersonation_chain,
+                    polling_interval_seconds=self.polling_interval_seconds,
+                    delete_on_error=self.delete_on_error,
+                ),
+                method_name="execute_complete",
+            )
         except AlreadyExists:
             if not self.use_if_exists:
                 raise
@@ -910,152 +915,13 @@ class DataprocCreateClusterOperator(GoogleCloudBaseOperator):
         cluster_state = event["cluster_state"]
         cluster_name = event["cluster_name"]
 
-        if cluster_state == ClusterStatus.State.ERROR:
+        if cluster_state == ClusterStatus.State(ClusterStatus.State.DELETING).name:
             raise AirflowException(f"Cluster is in ERROR state:\n{cluster_name}")
 
         self.log.info("%s completed successfully.", self.task_id)
         return event["cluster"]
 
 
-# TODO: Remove one day
-@deprecated(
-    planned_removal_date="March 01, 2025",
-    use_instead="DataprocUpdateClusterOperator",
-    category=AirflowProviderDeprecationWarning,
-)
-class DataprocScaleClusterOperator(GoogleCloudBaseOperator):
-    """
-    Scale, up or down, a cluster on Google Cloud Dataproc.
-
-    The operator will wait until the cluster is re-scaled.
-
-    Example usage:
-
-    .. code-block:: python
-
-        t1 = DataprocClusterScaleOperator(
-            task_id="dataproc_scale",
-            project_id="my-project",
-            cluster_name="cluster-1",
-            num_workers=10,
-            num_preemptible_workers=10,
-            graceful_decommission_timeout="1h",
-        )
-
-    .. seealso::
-        For more detail on about scaling clusters have a look at the reference:
-        https://cloud.google.com/dataproc/docs/concepts/configuring-clusters/scaling-clusters
-
-    :param cluster_name: The name of the cluster to scale. (templated)
-    :param project_id: The ID of the google cloud project in which
-        the cluster runs. (templated)
-    :param region: The region for the dataproc cluster. (templated)
-    :param num_workers: The new number of workers
-    :param num_preemptible_workers: The new number of preemptible workers
-    :param graceful_decommission_timeout: Timeout for graceful YARN decommissioning.
-        Maximum value is 1d
-    :param gcp_conn_id: The connection ID to use connecting to Google Cloud.
-    :param impersonation_chain: Optional service account to impersonate using short-term
-        credentials, or chained list of accounts required to get the access_token
-        of the last account in the list, which will be impersonated in the request.
-        If set as a string, the account must grant the originating account
-        the Service Account Token Creator IAM role.
-        If set as a sequence, the identities from the list must grant
-        Service Account Token Creator IAM role to the directly preceding identity, with first
-        account from the list granting this role to the originating account (templated).
-    """
-
-    template_fields: Sequence[str] = ("cluster_name", "project_id", "region", "impersonation_chain")
-
-    operator_extra_links = (DataprocLink(),)
-
-    def __init__(
-        self,
-        *,
-        cluster_name: str,
-        project_id: str = PROVIDE_PROJECT_ID,
-        region: str = "global",
-        num_workers: int = 2,
-        num_preemptible_workers: int = 0,
-        graceful_decommission_timeout: str | None = None,
-        gcp_conn_id: str = "google_cloud_default",
-        impersonation_chain: str | Sequence[str] | None = None,
-        **kwargs,
-    ) -> None:
-        super().__init__(**kwargs)
-        self.project_id = project_id
-        self.region = region
-        self.cluster_name = cluster_name
-        self.num_workers = num_workers
-        self.num_preemptible_workers = num_preemptible_workers
-        self.graceful_decommission_timeout = graceful_decommission_timeout
-        self.gcp_conn_id = gcp_conn_id
-        self.impersonation_chain = impersonation_chain
-
-    def _build_scale_cluster_data(self) -> dict:
-        scale_data = {
-            "config": {
-                "worker_config": {"num_instances": self.num_workers},
-                "secondary_worker_config": {"num_instances": self.num_preemptible_workers},
-            }
-        }
-        return scale_data
-
-    @property
-    def _graceful_decommission_timeout_object(self) -> dict[str, int] | None:
-        if not self.graceful_decommission_timeout:
-            return None
-
-        timeout = None
-        match = re.fullmatch(r"(\d+)([smdh])", self.graceful_decommission_timeout)
-        if match:
-            val = int(match.group(1))
-            unit = match.group(2)
-            if unit == "s":
-                timeout = val
-            elif unit == "m":
-                timeout = int(timedelta(minutes=val).total_seconds())
-            elif unit == "h":
-                timeout = int(timedelta(hours=val).total_seconds())
-            elif unit == "d":
-                timeout = int(timedelta(days=val).total_seconds())
-
-        if not timeout:
-            raise AirflowException(
-                "DataprocClusterScaleOperator "
-                " should be expressed in day, hours, minutes or seconds. "
-                " i.e. 1d, 4h, 10m, 30s"
-            )
-
-        return {"seconds": timeout}
-
-    def execute(self, context: Context) -> None:
-        """Scale, up or down, a cluster on Google Cloud Dataproc."""
-        self.log.info("Scaling cluster: %s", self.cluster_name)
-
-        scaling_cluster_data = self._build_scale_cluster_data()
-        update_mask = ["config.worker_config.num_instances", "config.secondary_worker_config.num_instances"]
-
-        hook = DataprocHook(gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain)
-        # Save data required to display extra link no matter what the cluster status will be
-        DataprocLink.persist(
-            context=context,
-            task_instance=self,
-            url=DATAPROC_CLUSTER_LINK_DEPRECATED,
-            resource=self.cluster_name,
-        )
-        operation = hook.update_cluster(
-            project_id=self.project_id,
-            region=self.region,
-            cluster_name=self.cluster_name,
-            cluster=scaling_cluster_data,
-            graceful_decommission_timeout=self._graceful_decommission_timeout_object,
-            update_mask={"paths": update_mask},
-        )
-        operation.result()
-        self.log.info("Cluster scaling finished")
-
-
 class DataprocDeleteClusterOperator(GoogleCloudBaseOperator):
     """
     Delete a cluster in a project.
@@ -1086,7 +952,13 @@ class DataprocDeleteClusterOperator(GoogleCloudBaseOperator):
     :param polling_interval_seconds: Time (seconds) to wait between calls to check the cluster status.
     """
 
-    template_fields: Sequence[str] = ("project_id", "region", "cluster_name", "impersonation_chain")
+    template_fields: Sequence[str] = (
+        "project_id",
+        "region",
+        "cluster_name",
+        "gcp_conn_id",
+        "impersonation_chain",
+    )
 
     def __init__(
         self,
@@ -1161,7 +1033,7 @@ class DataprocDeleteClusterOperator(GoogleCloudBaseOperator):
         """
         if event and event["status"] == "error":
             raise AirflowException(event["message"])
-        elif event is None:
+        if event is None:
             raise AirflowException("No event received in trigger callback")
         self.log.info("Cluster deleted.")
 
@@ -1212,6 +1084,7 @@ class _DataprocStartStopClusterBaseOperator(GoogleCloudBaseOperator):
         "region",
         "project_id",
         "request_id",
+        "gcp_conn_id",
         "impersonation_chain",
     )
 
@@ -1315,7 +1188,6 @@ class DataprocStartClusterOperator(_DataprocStartStopClusterBaseOperator):
         cluster = super().execute(context)
         DataprocClusterLink.persist(
             context=context,
-            operator=self,
             cluster_id=self.cluster_name,
             project_id=self._get_project_id(),
             region=self.region,
@@ -1461,8 +1333,7 @@ class DataprocJobBaseOperator(GoogleCloudBaseOperator):
         """Initialize `self.job_template` with default values."""
         if self.project_id is None:
             raise AirflowException(
-                "project id should either be set via project_id "
-                "parameter or retrieved from the connection,"
+                "project id should either be set via project_id parameter or retrieved from the connection,"
             )
         job_template = DataProcJobBuilder(
             project_id=self.project_id,
@@ -1497,7 +1368,11 @@ class DataprocJobBaseOperator(GoogleCloudBaseOperator):
             self.log.info("Job %s submitted successfully.", job_id)
             # Save data required for extra links no matter what the job status will be
             DataprocLink.persist(
-                context=context, task_instance=self, url=DATAPROC_JOB_LINK_DEPRECATED, resource=job_id
+                context=context,
+                url=DATAPROC_JOB_LINK_DEPRECATED,
+                resource=job_id,
+                region=self.region,
+                project_id=self.project_id,
             )
 
             if self.deferrable:
@@ -1517,8 +1392,7 @@ class DataprocJobBaseOperator(GoogleCloudBaseOperator):
                 self.hook.wait_for_job(job_id=job_id, region=self.region, project_id=self.project_id)
                 self.log.info("Job %s completed successfully.", job_id)
             return job_id
-        else:
-            raise AirflowException("Create a job template before")
+        raise AirflowException("Create a job template before")
 
     def execute_complete(self, context, event=None) -> None:
         """
@@ -1556,7 +1430,7 @@ class DataprocCreateWorkflowTemplateOperator(GoogleCloudBaseOperator):
     :param metadata: Additional metadata that is provided to the method.
     """
 
-    template_fields: Sequence[str] = ("region", "template")
+    template_fields: Sequence[str] = ("region", "template", "gcp_conn_id")
     template_fields_renderers = {"template": "json"}
     operator_extra_links = (DataprocWorkflowTemplateLink(),)
 
@@ -1602,7 +1476,6 @@ class DataprocCreateWorkflowTemplateOperator(GoogleCloudBaseOperator):
         if project_id:
             DataprocWorkflowTemplateLink.persist(
                 context=context,
-                operator=self,
                 workflow_template_id=self.template["id"],
                 region=self.region,
                 project_id=project_id,
@@ -1651,7 +1524,13 @@ class DataprocInstantiateWorkflowTemplateOperator(GoogleCloudBaseOperator):
     :param cancel_on_kill: Flag which indicates whether cancel the workflow, when on_kill is called
     """
 
-    template_fields: Sequence[str] = ("template_id", "impersonation_chain", "request_id", "parameters")
+    template_fields: Sequence[str] = (
+        "template_id",
+        "gcp_conn_id",
+        "impersonation_chain",
+        "request_id",
+        "parameters",
+    )
     template_fields_renderers = {"parameters": "json"}
     operator_extra_links = (DataprocWorkflowLink(),)
 
@@ -1714,7 +1593,6 @@ class DataprocInstantiateWorkflowTemplateOperator(GoogleCloudBaseOperator):
         if project_id:
             DataprocWorkflowLink.persist(
                 context=context,
-                operator=self,
                 workflow_id=workflow_id,
                 region=self.region,
                 project_id=project_id,
@@ -1800,7 +1678,7 @@ class DataprocInstantiateInlineWorkflowTemplateOperator(GoogleCloudBaseOperator)
     :param cancel_on_kill: Flag which indicates whether cancel the workflow, when on_kill is called
     """
 
-    template_fields: Sequence[str] = ("template", "impersonation_chain")
+    template_fields: Sequence[str] = ("template", "gcp_conn_id", "impersonation_chain")
     template_fields_renderers = {"template": "json"}
     operator_extra_links = (DataprocWorkflowLink(),)
 
@@ -1870,7 +1748,6 @@ class DataprocInstantiateInlineWorkflowTemplateOperator(GoogleCloudBaseOperator)
         if project_id:
             DataprocWorkflowLink.persist(
                 context=context,
-                operator=self,
                 workflow_id=workflow_id,
                 region=self.region,
                 project_id=project_id,
@@ -1969,7 +1846,14 @@ class DataprocSubmitJobOperator(GoogleCloudBaseOperator):
     :param wait_timeout: How many seconds wait for job to be ready. Used only if ``asynchronous`` is False
     """
 
-    template_fields: Sequence[str] = ("project_id", "region", "job", "impersonation_chain", "request_id")
+    template_fields: Sequence[str] = (
+        "project_id",
+        "region",
+        "job",
+        "gcp_conn_id",
+        "impersonation_chain",
+        "request_id",
+    )
     template_fields_renderers = {"job": "json"}
 
     operator_extra_links = (DataprocJobLink(),)
@@ -2044,7 +1928,6 @@ class DataprocSubmitJobOperator(GoogleCloudBaseOperator):
         if project_id:
             DataprocJobLink.persist(
                 context=context,
-                operator=self,
                 job_id=new_job_id,
                 region=self.region,
                 project_id=project_id,
@@ -2056,9 +1939,9 @@ class DataprocSubmitJobOperator(GoogleCloudBaseOperator):
             state = job.status.state
             if state == JobStatus.State.DONE:
                 return self.job_id
-            elif state == JobStatus.State.ERROR:
+            if state == JobStatus.State.ERROR:
                 raise AirflowException(f"Job failed:\n{job}")
-            elif state == JobStatus.State.CANCELLED:
+            if state == JobStatus.State.CANCELLED:
                 raise AirflowException(f"Job was cancelled:\n{job}")
             self.defer(
                 trigger=DataprocSubmitTrigger(
@@ -2090,9 +1973,9 @@ class DataprocSubmitJobOperator(GoogleCloudBaseOperator):
         job_state = event["job_state"]
         job_id = event["job_id"]
         job = event["job"]
-        if job_state == JobStatus.State.ERROR:
+        if job_state == JobStatus.State.ERROR.name:  # type: ignore
             raise AirflowException(f"Job {job_id} failed:\n{job}")
-        if job_state == JobStatus.State.CANCELLED:
+        if job_state == JobStatus.State.CANCELLED.name:  # type: ignore
             raise AirflowException(f"Job {job_id} was cancelled:\n{job}")
         self.log.info("%s completed successfully.", self.task_id)
         return job_id
@@ -2169,6 +2052,7 @@ class DataprocUpdateClusterOperator(GoogleCloudBaseOperator):
         "region",
         "request_id",
         "project_id",
+        "gcp_conn_id",
         "impersonation_chain",
     )
     operator_extra_links = (DataprocClusterLink(),)
@@ -2217,7 +2101,6 @@ class DataprocUpdateClusterOperator(GoogleCloudBaseOperator):
         if project_id:
             DataprocClusterLink.persist(
                 context=context,
-                operator=self,
                 cluster_id=self.cluster_name,
                 project_id=project_id,
                 region=self.region,
@@ -2305,6 +2188,7 @@ class DataprocDiagnoseClusterOperator(GoogleCloudBaseOperator):
         "project_id",
         "region",
         "cluster_name",
+        "gcp_conn_id",
         "impersonation_chain",
         "tarball_gcs_dir",
         "diagnosis_interval",
@@ -2451,6 +2335,7 @@ class DataprocCreateBatchOperator(GoogleCloudBaseOperator):
         "batch",
         "batch_id",
         "region",
+        "gcp_conn_id",
         "impersonation_chain",
     )
     operator_extra_links = (DataprocBatchLink(),)
@@ -2516,7 +2401,6 @@ class DataprocCreateBatchOperator(GoogleCloudBaseOperator):
             # Persist the link earlier so users can observe the progress
             DataprocBatchLink.persist(
                 context=context,
-                operator=self,
                 project_id=self.project_id,
                 region=self.region,
                 batch_id=self.batch_id,
@@ -2528,6 +2412,8 @@ class DataprocCreateBatchOperator(GoogleCloudBaseOperator):
             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,
@@ -2551,7 +2437,6 @@ class DataprocCreateBatchOperator(GoogleCloudBaseOperator):
 
         DataprocBatchLink.persist(
             context=context,
-            operator=self,
             project_id=self.project_id,
             region=self.region,
             batch_id=batch_id,
@@ -2601,7 +2486,7 @@ class DataprocCreateBatchOperator(GoogleCloudBaseOperator):
                 if not self.hook.check_error_for_resource_is_not_ready_msg(batch.state_message):
                     break
 
-        self.handle_batch_status(context, batch.state, batch_id, batch.state_message)
+        self.handle_batch_status(context, batch.state.name, batch_id, batch.state_message)
         return Batch.to_dict(batch)
 
     @cached_property
@@ -2626,21 +2511,21 @@ class DataprocCreateBatchOperator(GoogleCloudBaseOperator):
             self.operation.cancel()
 
     def handle_batch_status(
-        self, context: Context, state: Batch.State, batch_id: str, state_message: str | None = None
+        self, context: Context, state: str, batch_id: str, state_message: str | None = None
     ) -> None:
         # The existing batch may be a number of states other than 'SUCCEEDED'\
         # wait_for_operation doesn't fail if the job is cancelled, so we will check for it here which also
         # finds a cancelling|canceled|unspecified job from wait_for_batch or the deferred trigger
         link = DATAPROC_BATCH_LINK.format(region=self.region, project_id=self.project_id, batch_id=batch_id)
-        if state == Batch.State.FAILED:
+        if state == Batch.State.FAILED.name:  # type: ignore
             raise AirflowException(
-                f"Batch job {batch_id} failed with error: {state_message}\nDriver Logs: {link}"
+                f"Batch job {batch_id} failed with error: {state_message}.\nDriver logs: {link}"
             )
-        if state in (Batch.State.CANCELLED, Batch.State.CANCELLING):
-            raise AirflowException(f"Batch job {batch_id} was cancelled. Driver logs: {link}")
-        if state == Batch.State.STATE_UNSPECIFIED:
-            raise AirflowException(f"Batch job {batch_id} unspecified. Driver logs: {link}")
-        self.log.info("Batch job %s completed. Driver logs: %s", batch_id, link)
+        if state in (Batch.State.CANCELLED.name, Batch.State.CANCELLING.name):  # type: ignore
+            raise AirflowException(f"Batch job {batch_id} was cancelled.\nDriver logs: {link}")
+        if state == Batch.State.STATE_UNSPECIFIED.name:  # type: ignore
+            raise AirflowException(f"Batch job {batch_id} unspecified.\nDriver logs: {link}")
+        self.log.info("Batch job %s completed.\nDriver logs: %s", batch_id, link)
 
     def retry_batch_creation(
         self,
@@ -2708,6 +2593,31 @@ class DataprocCreateBatchOperator(GoogleCloudBaseOperator):
                 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,62}$")
+        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):
     """
@@ -2734,7 +2644,13 @@ class DataprocDeleteBatchOperator(GoogleCloudBaseOperator):
         account from the list granting this role to the originating account (templated).
     """
 
-    template_fields: Sequence[str] = ("batch_id", "region", "project_id", "impersonation_chain")
+    template_fields: Sequence[str] = (
+        "batch_id",
+        "region",
+        "project_id",
+        "gcp_conn_id",
+        "impersonation_chain",
+    )
 
     def __init__(
         self,
@@ -2798,7 +2714,13 @@ class DataprocGetBatchOperator(GoogleCloudBaseOperator):
         account from the list granting this role to the originating account (templated).
     """
 
-    template_fields: Sequence[str] = ("batch_id", "region", "project_id", "impersonation_chain")
+    template_fields: Sequence[str] = (
+        "batch_id",
+        "region",
+        "project_id",
+        "gcp_conn_id",
+        "impersonation_chain",
+    )
     operator_extra_links = (DataprocBatchLink(),)
 
     def __init__(
@@ -2839,7 +2761,6 @@ class DataprocGetBatchOperator(GoogleCloudBaseOperator):
         if project_id:
             DataprocBatchLink.persist(
                 context=context,
-                operator=self,
                 project_id=project_id,
                 region=self.region,
                 batch_id=self.batch_id,
@@ -2875,7 +2796,7 @@ class DataprocListBatchesOperator(GoogleCloudBaseOperator):
     :param order_by: How to order results as specified in ListBatchesRequest
     """
 
-    template_fields: Sequence[str] = ("region", "project_id", "impersonation_chain")
+    template_fields: Sequence[str] = ("region", "project_id", "gcp_conn_id", "impersonation_chain")
     operator_extra_links = (DataprocBatchesListLink(),)
 
     def __init__(
@@ -2922,7 +2843,7 @@ class DataprocListBatchesOperator(GoogleCloudBaseOperator):
         )
         project_id = self.project_id or hook.project_id
         if project_id:
-            DataprocBatchesListLink.persist(context=context, operator=self, project_id=project_id)
+            DataprocBatchesListLink.persist(context=context, project_id=project_id)
         return [Batch.to_dict(result) for result in results]
 
 
@@ -2949,7 +2870,13 @@ class DataprocCancelOperationOperator(GoogleCloudBaseOperator):
         account from the list granting this role to the originating account (templated).
     """
 
-    template_fields: Sequence[str] = ("operation_name", "region", "project_id", "impersonation_chain")
+    template_fields: Sequence[str] = (
+        "operation_name",
+        "region",
+        "project_id",
+        "gcp_conn_id",
+        "impersonation_chain",
+    )
 
     def __init__(
         self,