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

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

@@ -23,23 +23,26 @@ from collections.abc import Sequence
 from functools import cached_property
 from typing import TYPE_CHECKING
 
+from google.api_core.exceptions import AlreadyExists, NotFound
+from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
+from google.cloud.managedkafka_v1 import types
+
 from airflow.exceptions import AirflowException
 from airflow.providers.google.cloud.hooks.managed_kafka import ManagedKafkaHook
 from airflow.providers.google.cloud.links.managed_kafka import (
     ApacheKafkaClusterLink,
     ApacheKafkaClusterListLink,
+    ApacheKafkaConsumerGroupLink,
     ApacheKafkaTopicLink,
 )
 from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
-from google.api_core.exceptions import AlreadyExists, NotFound
-from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
-from google.cloud.managedkafka_v1 import types
 
 if TYPE_CHECKING:
-    from airflow.utils.context import Context
     from google.api_core.retry import Retry
     from google.protobuf.field_mask_pb2 import FieldMask
 
+    from airflow.utils.context import Context
+
 
 class ManagedKafkaBaseOperator(GoogleCloudBaseOperator):
     """
@@ -786,3 +789,267 @@ class ManagedKafkaDeleteTopicOperator(ManagedKafkaBaseOperator):
         except NotFound as not_found_err:
             self.log.info("The Apache Kafka topic ID %s does not exist.", self.topic_id)
             raise AirflowException(not_found_err)
+
+
+class ManagedKafkaListConsumerGroupsOperator(ManagedKafkaBaseOperator):
+    """
+    List the consumer groups in a given cluster.
+
+    :param project_id: Required. The ID of the Google Cloud project that the service belongs to.
+    :param location: Required. The ID of the Google Cloud region that the service belongs to.
+    :param cluster_id: Required. The ID of the cluster whose consumer groups are to be listed.
+    :param page_size: Optional. The maximum number of consumer groups to return. The service may return
+        fewer than this value. If unset or zero, all consumer groups for the parent is returned.
+    :param page_token: Optional. A page token, received from a previous ``ListConsumerGroups`` call.
+        Provide this to retrieve the subsequent page. When paginating, all other parameters provided to
+        ``ListConsumerGroups`` must match the call that provided the page token.
+    :param retry: Designation of what errors, if any, should be retried.
+    :param timeout: The timeout for this request.
+    :param metadata: Strings which should be sent along with the request as metadata.
+    :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] = tuple({"cluster_id"} | set(ManagedKafkaBaseOperator.template_fields))
+    operator_extra_links = (ApacheKafkaClusterLink(),)
+
+    def __init__(
+        self,
+        cluster_id: str,
+        page_size: int | None = None,
+        page_token: str | None = None,
+        *args,
+        **kwargs,
+    ) -> None:
+        super().__init__(*args, **kwargs)
+        self.cluster_id = cluster_id
+        self.page_size = page_size
+        self.page_token = page_token
+
+    def execute(self, context: Context):
+        ApacheKafkaClusterLink.persist(context=context, task_instance=self, cluster_id=self.cluster_id)
+        self.log.info("Listing Consumer Groups for cluster %s.", self.cluster_id)
+        try:
+            consumer_group_list_pager = self.hook.list_consumer_groups(
+                project_id=self.project_id,
+                location=self.location,
+                cluster_id=self.cluster_id,
+                page_size=self.page_size,
+                page_token=self.page_token,
+                retry=self.retry,
+                timeout=self.timeout,
+                metadata=self.metadata,
+            )
+            self.xcom_push(
+                context=context,
+                key="consumer_group_page",
+                value=types.ListConsumerGroupsResponse.to_dict(consumer_group_list_pager._response),
+            )
+        except Exception as error:
+            raise AirflowException(error)
+        return [types.ConsumerGroup.to_dict(consumer_group) for consumer_group in consumer_group_list_pager]
+
+
+class ManagedKafkaGetConsumerGroupOperator(ManagedKafkaBaseOperator):
+    """
+    Return the properties of a single consumer group.
+
+    :param project_id: Required. The ID of the Google Cloud project that the service belongs to.
+    :param location: Required. The ID of the Google Cloud region that the service belongs to.
+    :param cluster_id: Required. The ID of the cluster whose consumer group is to be returned.
+    :param consumer_group_id: Required. The ID of the consumer group whose configuration to return.
+    :param retry: Designation of what errors, if any, should be retried.
+    :param timeout: The timeout for this request.
+    :param metadata: Strings which should be sent along with the request as metadata.
+    :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] = tuple(
+        {"cluster_id", "consumer_group_id"} | set(ManagedKafkaBaseOperator.template_fields)
+    )
+    operator_extra_links = (ApacheKafkaConsumerGroupLink(),)
+
+    def __init__(
+        self,
+        cluster_id: str,
+        consumer_group_id: str,
+        *args,
+        **kwargs,
+    ) -> None:
+        super().__init__(*args, **kwargs)
+        self.cluster_id = cluster_id
+        self.consumer_group_id = consumer_group_id
+
+    def execute(self, context: Context):
+        ApacheKafkaConsumerGroupLink.persist(
+            context=context,
+            task_instance=self,
+            cluster_id=self.cluster_id,
+            consumer_group_id=self.consumer_group_id,
+        )
+        self.log.info("Getting Consumer Group: %s", self.consumer_group_id)
+        try:
+            consumer_group = self.hook.get_consumer_group(
+                project_id=self.project_id,
+                location=self.location,
+                cluster_id=self.cluster_id,
+                consumer_group_id=self.consumer_group_id,
+                retry=self.retry,
+                timeout=self.timeout,
+                metadata=self.metadata,
+            )
+            self.log.info(
+                "The consumer group %s from cluster %s was retrieved.",
+                self.consumer_group_id,
+                self.cluster_id,
+            )
+            return types.ConsumerGroup.to_dict(consumer_group)
+        except NotFound as not_found_err:
+            self.log.info("The Consumer Group %s does not exist.", self.consumer_group_id)
+            raise AirflowException(not_found_err)
+
+
+class ManagedKafkaUpdateConsumerGroupOperator(ManagedKafkaBaseOperator):
+    """
+    Update the properties of a single consumer group.
+
+    :param project_id: Required. The ID of the Google Cloud project that the service belongs to.
+    :param location: Required. The ID of the Google Cloud region that the service belongs to.
+    :param cluster_id: Required. The ID of the cluster whose topic is to be updated.
+    :param consumer_group_id: Required. The ID of the consumer group whose configuration to update.
+    :param consumer_group: Required. The consumer_group to update. Its ``name`` field must be populated.
+    :param update_mask: Required. Field mask is used to specify the fields to be overwritten in the
+        ConsumerGroup resource by the update. The fields specified in the update_mask are relative to the
+        resource, not the full request. A field will be overwritten if it is in the mask.
+    :param retry: Designation of what errors, if any, should be retried.
+    :param timeout: The timeout for this request.
+    :param metadata: Strings which should be sent along with the request as metadata.
+    :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] = tuple(
+        {"cluster_id", "consumer_group_id", "consumer_group", "update_mask"}
+        | set(ManagedKafkaBaseOperator.template_fields)
+    )
+    operator_extra_links = (ApacheKafkaConsumerGroupLink(),)
+
+    def __init__(
+        self,
+        cluster_id: str,
+        consumer_group_id: str,
+        consumer_group: types.Topic | dict,
+        update_mask: FieldMask | dict,
+        *args,
+        **kwargs,
+    ) -> None:
+        super().__init__(*args, **kwargs)
+        self.cluster_id = cluster_id
+        self.consumer_group_id = consumer_group_id
+        self.consumer_group = consumer_group
+        self.update_mask = update_mask
+
+    def execute(self, context: Context):
+        ApacheKafkaConsumerGroupLink.persist(
+            context=context,
+            task_instance=self,
+            cluster_id=self.cluster_id,
+            consumer_group_id=self.consumer_group_id,
+        )
+        self.log.info("Updating an Apache Kafka consumer group.")
+        try:
+            consumer_group_obj = self.hook.update_consumer_group(
+                project_id=self.project_id,
+                location=self.location,
+                cluster_id=self.cluster_id,
+                consumer_group_id=self.consumer_group_id,
+                consumer_group=self.consumer_group,
+                update_mask=self.update_mask,
+                retry=self.retry,
+                timeout=self.timeout,
+                metadata=self.metadata,
+            )
+            self.log.info("Apache Kafka consumer group %s was updated.", self.consumer_group_id)
+            return types.ConsumerGroup.to_dict(consumer_group_obj)
+        except NotFound as not_found_err:
+            self.log.info("The Consumer Group %s does not exist.", self.consumer_group_id)
+            raise AirflowException(not_found_err)
+        except Exception as error:
+            raise AirflowException(error)
+
+
+class ManagedKafkaDeleteConsumerGroupOperator(ManagedKafkaBaseOperator):
+    """
+    Delete a single consumer group.
+
+    :param project_id: Required. The ID of the Google Cloud project that the service belongs to.
+    :param location: Required. The ID of the Google Cloud region that the service belongs to.
+    :param cluster_id: Required. The ID of the cluster whose consumer group is to be deleted.
+    :param consumer_group_id: Required. The ID of the consumer group to delete.
+    :param retry: Designation of what errors, if any, should be retried.
+    :param timeout: The timeout for this request.
+    :param metadata: Strings which should be sent along with the request as metadata.
+    :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] = tuple(
+        {"cluster_id", "consumer_group_id"} | set(ManagedKafkaBaseOperator.template_fields)
+    )
+
+    def __init__(
+        self,
+        cluster_id: str,
+        consumer_group_id: str,
+        *args,
+        **kwargs,
+    ) -> None:
+        super().__init__(*args, **kwargs)
+        self.cluster_id = cluster_id
+        self.consumer_group_id = consumer_group_id
+
+    def execute(self, context: Context):
+        try:
+            self.log.info("Deleting Apache Kafka consumer group: %s", self.consumer_group_id)
+            self.hook.delete_consumer_group(
+                project_id=self.project_id,
+                location=self.location,
+                cluster_id=self.cluster_id,
+                consumer_group_id=self.consumer_group_id,
+                retry=self.retry,
+                timeout=self.timeout,
+                metadata=self.metadata,
+            )
+            self.log.info("Apache Kafka consumer group was deleted.")
+        except NotFound as not_found_err:
+            self.log.info("The Apache Kafka consumer group ID %s does not exist.", self.consumer_group_id)
+            raise AirflowException(not_found_err)

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

@@ -22,16 +22,18 @@ from __future__ import annotations
 from collections.abc import Sequence
 from typing import TYPE_CHECKING
 
-from airflow.providers.google.cloud.hooks.natural_language import CloudNaturalLanguageHook
-from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
 from google.protobuf.json_format import MessageToDict
 
+from airflow.providers.google.cloud.hooks.natural_language import CloudNaturalLanguageHook
+from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
+
 if TYPE_CHECKING:
-    from airflow.utils.context import Context
     from google.api_core.retry import Retry
     from google.cloud.language_v1.types import Document, EncodingType
 
+    from airflow.utils.context import Context
+
 
 MetaData = Sequence[tuple[str, str]]
 

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

@@ -28,13 +28,6 @@ from __future__ import annotations
 from collections.abc import Sequence
 from typing import TYPE_CHECKING, Any, Callable
 
-from airflow.exceptions import AirflowException
-from airflow.providers.google.cloud.hooks.pubsub import PubSubHook
-from airflow.providers.google.cloud.links.pubsub import PubSubSubscriptionLink, PubSubTopicLink
-from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
-from airflow.providers.google.cloud.triggers.pubsub import PubsubPullTrigger
-from airflow.providers.google.common.consts import GOOGLE_DEFAULT_DEFERRABLE_METHOD_NAME
-from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
 from google.cloud.pubsub_v1.types import (
     DeadLetterPolicy,
@@ -47,10 +40,19 @@ from google.cloud.pubsub_v1.types import (
     SchemaSettings,
 )
 
+from airflow.exceptions import AirflowException
+from airflow.providers.google.cloud.hooks.pubsub import PubSubHook
+from airflow.providers.google.cloud.links.pubsub import PubSubSubscriptionLink, PubSubTopicLink
+from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
+from airflow.providers.google.cloud.triggers.pubsub import PubsubPullTrigger
+from airflow.providers.google.common.consts import GOOGLE_DEFAULT_DEFERRABLE_METHOD_NAME
+from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
+
 if TYPE_CHECKING:
-    from airflow.utils.context import Context
     from google.api_core.retry import Retry
 
+    from airflow.utils.context import Context
+
 
 class PubSubCreateTopicOperator(GoogleCloudBaseOperator):
     """

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

@@ -22,19 +22,21 @@ from __future__ import annotations
 from collections.abc import Sequence
 from typing import TYPE_CHECKING
 
+from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
+from google.protobuf.json_format import MessageToDict
+
 from airflow.exceptions import AirflowException
 from airflow.providers.google.cloud.hooks.speech_to_text import CloudSpeechToTextHook, RecognitionAudio
 from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
 from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
 from airflow.providers.google.common.links.storage import FileDetailsLink
-from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
-from google.protobuf.json_format import MessageToDict
 
 if TYPE_CHECKING:
-    from airflow.utils.context import Context
     from google.api_core.retry import Retry
     from google.cloud.speech_v1.types import RecognitionConfig
 
+    from airflow.utils.context import Context
+
 
 class CloudSpeechToTextRecognizeSpeechOperator(GoogleCloudBaseOperator):
     """

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

@@ -20,6 +20,9 @@ from __future__ import annotations
 from collections.abc import Sequence
 from typing import TYPE_CHECKING
 
+from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
+from google.cloud.monitoring_v3 import AlertPolicy, NotificationChannel
+
 from airflow.providers.google.cloud.hooks.stackdriver import StackdriverHook
 from airflow.providers.google.cloud.links.stackdriver import (
     StackdriverNotificationsLink,
@@ -27,13 +30,12 @@ from airflow.providers.google.cloud.links.stackdriver import (
 )
 from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
 from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
-from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
-from google.cloud.monitoring_v3 import AlertPolicy, NotificationChannel
 
 if TYPE_CHECKING:
-    from airflow.utils.context import Context
     from google.api_core.retry import Retry
 
+    from airflow.utils.context import Context
+
 
 class StackdriverListAlertPoliciesOperator(GoogleCloudBaseOperator):
     """

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

@@ -22,19 +22,21 @@ from __future__ import annotations
 from collections.abc import Sequence
 from typing import TYPE_CHECKING
 
+from google.api_core.exceptions import AlreadyExists
+from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
+from google.cloud.tasks_v2.types import Queue, Task
+
 from airflow.providers.google.cloud.hooks.tasks import CloudTasksHook
 from airflow.providers.google.cloud.links.cloud_tasks import CloudTasksLink, CloudTasksQueueLink
 from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
 from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
-from google.api_core.exceptions import AlreadyExists
-from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
-from google.cloud.tasks_v2.types import Queue, Task
 
 if TYPE_CHECKING:
-    from airflow.utils.context import Context
     from google.api_core.retry import Retry
     from google.protobuf.field_mask_pb2 import FieldMask
 
+    from airflow.utils.context import Context
+
 
 MetaData = Sequence[tuple[str, str]]
 

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

@@ -23,19 +23,21 @@ from collections.abc import Sequence
 from tempfile import NamedTemporaryFile
 from typing import TYPE_CHECKING
 
+from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
+
 from airflow.exceptions import AirflowException
 from airflow.providers.google.cloud.hooks.gcs import GCSHook
 from airflow.providers.google.cloud.hooks.text_to_speech import CloudTextToSpeechHook
 from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
 from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
 from airflow.providers.google.common.links.storage import FileDetailsLink
-from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
 
 if TYPE_CHECKING:
-    from airflow.utils.context import Context
     from google.api_core.retry import Retry
     from google.cloud.texttospeech_v1.types import AudioConfig, SynthesisInput, VoiceSelectionParams
 
+    from airflow.utils.context import Context
+
 
 class CloudTextToSpeechSynthesizeOperator(GoogleCloudBaseOperator):
     """

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

@@ -22,6 +22,9 @@ from __future__ import annotations
 from collections.abc import MutableMapping, MutableSequence, Sequence
 from typing import TYPE_CHECKING, cast
 
+from google.api_core.exceptions import GoogleAPICallError
+from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
+
 from airflow.exceptions import AirflowException
 from airflow.providers.google.cloud.hooks.translate import CloudTranslateHook, TranslateHook
 from airflow.providers.google.cloud.links.translate import (
@@ -35,11 +38,8 @@ from airflow.providers.google.cloud.links.translate import (
 )
 from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
 from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
-from google.api_core.exceptions import GoogleAPICallError
-from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
 
 if TYPE_CHECKING:
-    from airflow.utils.context import Context
     from google.api_core.retry import Retry
     from google.cloud.translate_v3.types import (
         BatchDocumentInputConfig,
@@ -55,6 +55,8 @@ if TYPE_CHECKING:
     )
     from google.cloud.translate_v3.types.translation_service import Glossary, GlossaryInputConfig
 
+    from airflow.utils.context import Context
+
 
 class CloudTranslateTextOperator(GoogleCloudBaseOperator):
     """

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

@@ -22,18 +22,20 @@ from __future__ import annotations
 from collections.abc import Sequence
 from typing import TYPE_CHECKING
 
+from google.protobuf.json_format import MessageToDict
+
 from airflow.exceptions import AirflowException
 from airflow.providers.google.cloud.hooks.speech_to_text import CloudSpeechToTextHook
 from airflow.providers.google.cloud.hooks.translate import CloudTranslateHook
 from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
 from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
 from airflow.providers.google.common.links.storage import FileDetailsLink
-from google.protobuf.json_format import MessageToDict
 
 if TYPE_CHECKING:
-    from airflow.utils.context import Context
     from google.cloud.speech_v1.types import RecognitionAudio, RecognitionConfig
 
+    from airflow.utils.context import Context
+
 
 class CloudTranslateSpeechOperator(GoogleCloudBaseOperator):
     """

airflow/providers/google/cloud/operators/vertex_ai/auto_ml.py

@@ -23,6 +23,12 @@ from __future__ import annotations
 from collections.abc import Sequence
 from typing import TYPE_CHECKING
 
+from google.api_core.exceptions import NotFound
+from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
+from google.cloud.aiplatform import datasets
+from google.cloud.aiplatform.models import Model
+from google.cloud.aiplatform_v1.types.training_pipeline import TrainingPipeline
+
 from airflow.exceptions import AirflowProviderDeprecationWarning
 from airflow.providers.google.cloud.hooks.vertex_ai.auto_ml import AutoMLHook
 from airflow.providers.google.cloud.links.vertex_ai import (
@@ -32,16 +38,12 @@ from airflow.providers.google.cloud.links.vertex_ai import (
 )
 from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
 from airflow.providers.google.common.deprecated import deprecated
-from google.api_core.exceptions import NotFound
-from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
-from google.cloud.aiplatform import datasets
-from google.cloud.aiplatform.models import Model
-from google.cloud.aiplatform_v1.types.training_pipeline import TrainingPipeline
 
 if TYPE_CHECKING:
-    from airflow.utils.context import Context
     from google.api_core.retry import Retry
 
+    from airflow.utils.context import Context
+
 
 class AutoMLTrainingJobBaseOperator(GoogleCloudBaseOperator):
     """The base class for operators that launch AutoML jobs on VertexAI."""

airflow/providers/google/cloud/operators/vertex_ai/batch_prediction_job.py

@@ -24,6 +24,10 @@ from collections.abc import Sequence
 from functools import cached_property
 from typing import TYPE_CHECKING, Any
 
+from google.api_core.exceptions import NotFound
+from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
+from google.cloud.aiplatform_v1.types import BatchPredictionJob
+
 from airflow.configuration import conf
 from airflow.exceptions import AirflowException
 from airflow.providers.google.cloud.hooks.vertex_ai.batch_prediction_job import BatchPredictionJobHook
@@ -33,15 +37,13 @@ from airflow.providers.google.cloud.links.vertex_ai import (
 )
 from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
 from airflow.providers.google.cloud.triggers.vertex_ai import CreateBatchPredictionJobTrigger
-from google.api_core.exceptions import NotFound
-from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
-from google.cloud.aiplatform_v1.types import BatchPredictionJob
 
 if TYPE_CHECKING:
-    from airflow.utils.context import Context
     from google.api_core.retry import Retry
     from google.cloud.aiplatform import BatchPredictionJob as BatchPredictionJobObject, Model, explain
 
+    from airflow.utils.context import Context
+
 
 class CreateBatchPredictionJobOperator(GoogleCloudBaseOperator):
     """

airflow/providers/google/cloud/operators/vertex_ai/custom_job.py

@@ -23,6 +23,12 @@ from collections.abc import Sequence
 from functools import cached_property
 from typing import TYPE_CHECKING, Any
 
+from google.api_core.exceptions import NotFound
+from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
+from google.cloud.aiplatform.models import Model
+from google.cloud.aiplatform_v1.types.dataset import Dataset
+from google.cloud.aiplatform_v1.types.training_pipeline import TrainingPipeline
+
 from airflow.configuration import conf
 from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
 from airflow.providers.google.cloud.hooks.vertex_ai.custom_job import CustomJobHook
@@ -38,14 +44,8 @@ from airflow.providers.google.cloud.triggers.vertex_ai import (
     CustomTrainingJobTrigger,
 )
 from airflow.providers.google.common.deprecated import deprecated
-from google.api_core.exceptions import NotFound
-from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
-from google.cloud.aiplatform.models import Model
-from google.cloud.aiplatform_v1.types.dataset import Dataset
-from google.cloud.aiplatform_v1.types.training_pipeline import TrainingPipeline
 
 if TYPE_CHECKING:
-    from airflow.utils.context import Context
     from google.api_core.retry import Retry
     from google.cloud.aiplatform import (
         CustomContainerTrainingJob,
@@ -53,6 +53,8 @@ if TYPE_CHECKING:
         CustomTrainingJob,
     )
 
+    from airflow.utils.context import Context
+
 
 class CustomTrainingJobBaseOperator(GoogleCloudBaseOperator):
     """The base class for operators that launch Custom jobs on VertexAI."""

airflow/providers/google/cloud/operators/vertex_ai/dataset.py

@@ -22,18 +22,20 @@ from __future__ import annotations
 from collections.abc import Sequence
 from typing import TYPE_CHECKING
 
-from airflow.providers.google.cloud.hooks.vertex_ai.dataset import DatasetHook
-from airflow.providers.google.cloud.links.vertex_ai import VertexAIDatasetLink, VertexAIDatasetListLink
-from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
 from google.api_core.exceptions import NotFound
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
 from google.cloud.aiplatform_v1.types import Dataset, ExportDataConfig, ImportDataConfig
 
+from airflow.providers.google.cloud.hooks.vertex_ai.dataset import DatasetHook
+from airflow.providers.google.cloud.links.vertex_ai import VertexAIDatasetLink, VertexAIDatasetListLink
+from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
+
 if TYPE_CHECKING:
-    from airflow.utils.context import Context
     from google.api_core.retry import Retry
     from google.protobuf.field_mask_pb2 import FieldMask
 
+    from airflow.utils.context import Context
+
 
 class CreateDatasetOperator(GoogleCloudBaseOperator):
     """

airflow/providers/google/cloud/operators/vertex_ai/endpoint_service.py

@@ -23,6 +23,10 @@ from __future__ import annotations
 from collections.abc import Sequence
 from typing import TYPE_CHECKING
 
+from google.api_core.exceptions import NotFound
+from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
+from google.cloud.aiplatform_v1.types import DeployedModel, Endpoint, endpoint_service
+
 from airflow.providers.google.cloud.hooks.vertex_ai.endpoint_service import EndpointServiceHook
 from airflow.providers.google.cloud.links.vertex_ai import (
     VertexAIEndpointLink,
@@ -30,15 +34,13 @@ from airflow.providers.google.cloud.links.vertex_ai import (
     VertexAIModelLink,
 )
 from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
-from google.api_core.exceptions import NotFound
-from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
-from google.cloud.aiplatform_v1.types import DeployedModel, Endpoint, endpoint_service
 
 if TYPE_CHECKING:
-    from airflow.utils.context import Context
     from google.api_core.retry import Retry
     from google.protobuf.field_mask_pb2 import FieldMask
 
+    from airflow.utils.context import Context
+
 
 class CreateEndpointOperator(GoogleCloudBaseOperator):
     """