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

airflow/providers/google/cloud/hooks/os_login.py

@@ -27,11 +27,12 @@ from __future__ import annotations
 from collections.abc import Sequence
 from typing import TYPE_CHECKING
 
-from airflow.providers.google.common.consts import CLIENT_INFO
-from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID, GoogleBaseHook
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
 from google.cloud.oslogin_v1 import ImportSshPublicKeyResponse, OsLoginServiceClient
 
+from airflow.providers.google.common.consts import CLIENT_INFO
+from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID, GoogleBaseHook
+
 if TYPE_CHECKING:
     from google.api_core.retry import Retry
 

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

@@ -33,6 +33,12 @@ from functools import cached_property
 from typing import TYPE_CHECKING, Any
 from uuid import uuid4
 
+from google.api_core.exceptions import AlreadyExists, GoogleAPICallError
+from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
+from google.cloud.exceptions import NotFound
+from google.cloud.pubsub_v1 import PublisherClient, SubscriberClient
+from google.cloud.pubsub_v1.types import PublisherOptions
+from google.pubsub_v1.services.subscriber.async_client import SubscriberAsyncClient
 from googleapiclient.errors import HttpError
 
 from airflow.exceptions import AirflowProviderDeprecationWarning
@@ -43,12 +49,6 @@ from airflow.providers.google.common.hooks.base_google import (
     GoogleBaseHook,
 )
 from airflow.version import version
-from google.api_core.exceptions import AlreadyExists, GoogleAPICallError
-from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
-from google.cloud.exceptions import NotFound
-from google.cloud.pubsub_v1 import PublisherClient, SubscriberClient
-from google.cloud.pubsub_v1.types import PublisherOptions
-from google.pubsub_v1.services.subscriber.async_client import SubscriberAsyncClient
 
 if TYPE_CHECKING:
     from google.api_core.retry import Retry

airflow/providers/google/cloud/hooks/secret_manager.py

@@ -23,8 +23,6 @@ from collections.abc import Sequence
 from functools import cached_property
 from typing import TYPE_CHECKING
 
-from airflow.providers.google.common.consts import CLIENT_INFO
-from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
 from google.cloud.secretmanager_v1 import (
     AccessSecretVersionResponse,
@@ -34,6 +32,9 @@ from google.cloud.secretmanager_v1 import (
     SecretVersion,
 )
 
+from airflow.providers.google.common.consts import CLIENT_INFO
+from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
+
 if TYPE_CHECKING:
     from google.api_core.retry import Retry
     from google.cloud.secretmanager_v1.services.secret_manager_service.pagers import ListSecretsPager

airflow/providers/google/cloud/hooks/spanner.py

@@ -22,14 +22,14 @@ from __future__ import annotations
 from collections.abc import Sequence
 from typing import TYPE_CHECKING, Callable, NamedTuple
 
+from google.api_core.exceptions import AlreadyExists, GoogleAPICallError
+from google.cloud.spanner_v1.client import Client
 from sqlalchemy import create_engine
 
 from airflow.exceptions import AirflowException
 from airflow.providers.common.sql.hooks.sql import DbApiHook
 from airflow.providers.google.common.consts import CLIENT_INFO
 from airflow.providers.google.common.hooks.base_google import GoogleBaseHook, get_field
-from google.api_core.exceptions import AlreadyExists, GoogleAPICallError
-from google.cloud.spanner_v1.client import Client
 
 if TYPE_CHECKING:
     from google.cloud.spanner_v1.database import Database

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

@@ -22,12 +22,13 @@ from __future__ import annotations
 from collections.abc import Sequence
 from typing import TYPE_CHECKING
 
-from airflow.providers.google.common.consts import CLIENT_INFO
-from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
 from google.cloud.speech_v1 import SpeechClient
 from google.cloud.speech_v1.types import RecognitionAudio, RecognitionConfig
 
+from airflow.providers.google.common.consts import CLIENT_INFO
+from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
+
 if TYPE_CHECKING:
     from google.api_core.retry import Retry
 

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

@@ -24,15 +24,15 @@ import json
 from collections.abc import Sequence
 from typing import TYPE_CHECKING, Any
 
-from googleapiclient.errors import HttpError
-
-from airflow.exceptions import AirflowException
-from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID, GoogleBaseHook
 from google.api_core.exceptions import InvalidArgument
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
 from google.cloud import monitoring_v3
 from google.cloud.monitoring_v3 import AlertPolicy, NotificationChannel
 from google.protobuf.field_mask_pb2 import FieldMask
+from googleapiclient.errors import HttpError
+
+from airflow.exceptions import AirflowException
+from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID, GoogleBaseHook
 
 if TYPE_CHECKING:
     from google.api_core.retry import Retry

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

@@ -22,13 +22,14 @@ from __future__ import annotations
 from collections.abc import Sequence
 from typing import TYPE_CHECKING
 
-from airflow.exceptions import AirflowException
-from airflow.providers.google.common.consts import CLIENT_INFO
-from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID, GoogleBaseHook
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
 from google.cloud.tasks_v2 import CloudTasksClient
 from google.cloud.tasks_v2.types import Queue, Task
 
+from airflow.exceptions import AirflowException
+from airflow.providers.google.common.consts import CLIENT_INFO
+from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID, GoogleBaseHook
+
 if TYPE_CHECKING:
     from google.api_core.retry import Retry
     from google.protobuf.field_mask_pb2 import FieldMask

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

@@ -22,8 +22,6 @@ from __future__ import annotations
 from collections.abc import Sequence
 from typing import TYPE_CHECKING
 
-from airflow.providers.google.common.consts import CLIENT_INFO
-from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
 from google.cloud.texttospeech_v1 import TextToSpeechClient
 from google.cloud.texttospeech_v1.types import (
@@ -33,6 +31,9 @@ from google.cloud.texttospeech_v1.types import (
     VoiceSelectionParams,
 )
 
+from airflow.providers.google.common.consts import CLIENT_INFO
+from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
+
 if TYPE_CHECKING:
     from google.api_core.retry import Retry
 

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

@@ -25,9 +25,6 @@ from typing import (
     cast,
 )
 
-from airflow.exceptions import AirflowException
-from airflow.providers.google.common.consts import CLIENT_INFO
-from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID, GoogleBaseHook
 from google.api_core.exceptions import GoogleAPICallError
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
 from google.api_core.retry import Retry
@@ -35,9 +32,12 @@ from google.cloud.translate_v2 import Client
 from google.cloud.translate_v3 import TranslationServiceClient
 from google.cloud.translate_v3.types.translation_service import GlossaryInputConfig
 
-if TYPE_CHECKING:
-    from proto import Message
+from airflow.exceptions import AirflowException
+from airflow.providers.google.common.consts import CLIENT_INFO
+from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID, GoogleBaseHook
+from airflow.providers.google.common.hooks.operation_helpers import OperationHelper
 
+if TYPE_CHECKING:
     from google.api_core.operation import Operation
     from google.cloud.translate_v3.services.translation_service import pagers
     from google.cloud.translate_v3.types import (
@@ -155,7 +155,7 @@ class CloudTranslateHook(GoogleBaseHook):
         )
 
 
-class TranslateHook(GoogleBaseHook):
+class TranslateHook(GoogleBaseHook, OperationHelper):
     """
     Hook for Google Cloud translation (Advanced) using client version V3.
 
@@ -221,15 +221,6 @@ class TranslateHook(GoogleBaseHook):
             error = operation.exception(timeout=timeout)
             raise AirflowException(error)
 
-    @staticmethod
-    def wait_for_operation_result(operation: Operation, timeout: int | None = None) -> Message:
-        """Wait for long-lasting operation to complete."""
-        try:
-            return operation.result(timeout=timeout)
-        except GoogleAPICallError:
-            error = operation.exception(timeout=timeout)
-            raise AirflowException(error)
-
     @staticmethod
     def extract_object_id(obj: dict) -> str:
         """Return unique id of the object."""

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

@@ -23,9 +23,6 @@ import warnings
 from collections.abc import Sequence
 from typing import TYPE_CHECKING
 
-from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
-from airflow.providers.google.common.deprecated import deprecated
-from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
 from google.api_core.client_options import ClientOptions
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
 from google.cloud.aiplatform import (
@@ -39,6 +36,11 @@ from google.cloud.aiplatform import (
 )
 from google.cloud.aiplatform_v1 import JobServiceClient, PipelineServiceClient
 
+from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
+from airflow.providers.google.common.deprecated import deprecated
+from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
+from airflow.providers.google.common.hooks.operation_helpers import OperationHelper
+
 if TYPE_CHECKING:
     from google.api_core.operation import Operation
     from google.api_core.retry import Retry
@@ -46,7 +48,7 @@ if TYPE_CHECKING:
     from google.cloud.aiplatform_v1.types import TrainingPipeline
 
 
-class AutoMLHook(GoogleBaseHook):
+class AutoMLHook(GoogleBaseHook, OperationHelper):
     """Hook for Google Cloud Vertex AI Auto ML APIs."""
 
     def __init__(
@@ -252,14 +254,6 @@ class AutoMLHook(GoogleBaseHook):
         """Return unique id of the Training pipeline."""
         return resource_name.rpartition("/")[-1]
 
-    def wait_for_operation(self, operation: Operation, timeout: float | None = None):
-        """Wait for long-lasting operation to complete."""
-        try:
-            return operation.result(timeout=timeout)
-        except Exception:
-            error = operation.exception(timeout=timeout)
-            raise AirflowException(error)
-
     def cancel_auto_ml_job(self) -> None:
         """Cancel Auto ML Job for training pipeline."""
         if self._job:

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

@@ -23,21 +23,23 @@ import asyncio
 from collections.abc import Sequence
 from typing import TYPE_CHECKING
 
-from airflow.exceptions import AirflowException
-from airflow.providers.google.common.consts import CLIENT_INFO
-from airflow.providers.google.common.hooks.base_google import GoogleBaseAsyncHook, GoogleBaseHook
 from google.api_core.client_options import ClientOptions
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
 from google.cloud.aiplatform import BatchPredictionJob, Model, explain
 from google.cloud.aiplatform_v1 import JobServiceAsyncClient, JobServiceClient, JobState, types
 
+from airflow.exceptions import AirflowException
+from airflow.providers.google.common.consts import CLIENT_INFO
+from airflow.providers.google.common.hooks.base_google import GoogleBaseAsyncHook, GoogleBaseHook
+from airflow.providers.google.common.hooks.operation_helpers import OperationHelper
+
 if TYPE_CHECKING:
     from google.api_core.operation import Operation
     from google.api_core.retry import AsyncRetry, Retry
     from google.cloud.aiplatform_v1.services.job_service.pagers import ListBatchPredictionJobsPager
 
 
-class BatchPredictionJobHook(GoogleBaseHook):
+class BatchPredictionJobHook(GoogleBaseHook, OperationHelper):
     """Hook for Google Cloud Vertex AI Batch Prediction Job APIs."""
 
     def __init__(
@@ -64,14 +66,6 @@ class BatchPredictionJobHook(GoogleBaseHook):
             credentials=self.get_credentials(), client_info=self.client_info, client_options=client_options
         )
 
-    def wait_for_operation(self, operation: Operation, timeout: float | None = None):
-        """Wait for long-lasting operation to complete."""
-        try:
-            return operation.result(timeout=timeout)
-        except Exception:
-            error = operation.exception(timeout=timeout)
-            raise AirflowException(error)
-
     @staticmethod
     def extract_batch_prediction_job_id(obj: dict) -> str:
         """Return unique id of the batch_prediction_job."""

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

@@ -23,10 +23,6 @@ import asyncio
 from collections.abc import Sequence
 from typing import TYPE_CHECKING, Any
 
-from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
-from airflow.providers.google.common.consts import CLIENT_INFO
-from airflow.providers.google.common.deprecated import deprecated
-from airflow.providers.google.common.hooks.base_google import GoogleBaseAsyncHook, GoogleBaseHook
 from google.api_core.client_options import ClientOptions
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
 from google.cloud.aiplatform import (
@@ -46,6 +42,12 @@ from google.cloud.aiplatform_v1 import (
     types,
 )
 
+from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
+from airflow.providers.google.common.consts import CLIENT_INFO
+from airflow.providers.google.common.deprecated import deprecated
+from airflow.providers.google.common.hooks.base_google import GoogleBaseAsyncHook, GoogleBaseHook
+from airflow.providers.google.common.hooks.operation_helpers import OperationHelper
+
 if TYPE_CHECKING:
     from google.api_core.operation import Operation
     from google.api_core.retry import AsyncRetry, Retry
@@ -58,7 +60,7 @@ if TYPE_CHECKING:
     from google.cloud.aiplatform_v1.types import CustomJob, PipelineJob, TrainingPipeline
 
 
-class CustomJobHook(GoogleBaseHook):
+class CustomJobHook(GoogleBaseHook, OperationHelper):
     """Hook for Google Cloud Vertex AI Custom Job APIs."""
 
     def __init__(
@@ -276,14 +278,6 @@ class CustomJobHook(GoogleBaseHook):
         """Return a unique Custom Job id from a serialized TrainingPipeline proto."""
         return training_pipeline["training_task_metadata"]["backingCustomJob"].rpartition("/")[-1]
 
-    def wait_for_operation(self, operation: Operation, timeout: float | None = None):
-        """Wait for long-lasting operation to complete."""
-        try:
-            return operation.result(timeout=timeout)
-        except Exception:
-            error = operation.exception(timeout=timeout)
-            raise AirflowException(error)
-
     def cancel_job(self) -> None:
         """Cancel Job for training pipeline."""
         if self._job:

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

@@ -22,13 +22,14 @@ from __future__ import annotations
 from collections.abc import Sequence
 from typing import TYPE_CHECKING
 
-from airflow.exceptions import AirflowException
-from airflow.providers.google.common.consts import CLIENT_INFO
-from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
 from google.api_core.client_options import ClientOptions
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
 from google.cloud.aiplatform_v1 import DatasetServiceClient
 
+from airflow.providers.google.common.consts import CLIENT_INFO
+from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
+from airflow.providers.google.common.hooks.operation_helpers import OperationHelper
+
 if TYPE_CHECKING:
     from google.api_core.operation import Operation
     from google.api_core.retry import Retry
@@ -41,7 +42,7 @@ if TYPE_CHECKING:
     from google.protobuf.field_mask_pb2 import FieldMask
 
 
-class DatasetHook(GoogleBaseHook):
+class DatasetHook(GoogleBaseHook, OperationHelper):
     """Hook for Google Cloud Vertex AI Dataset APIs."""
 
     def get_dataset_service_client(self, region: str | None = None) -> DatasetServiceClient:
@@ -55,14 +56,6 @@ class DatasetHook(GoogleBaseHook):
             credentials=self.get_credentials(), client_info=CLIENT_INFO, client_options=client_options
         )
 
-    def wait_for_operation(self, operation: Operation, timeout: float | None = None):
-        """Wait for long-lasting operation to complete."""
-        try:
-            return operation.result(timeout=timeout)
-        except Exception:
-            error = operation.exception(timeout=timeout)
-            raise AirflowException(error)
-
     @staticmethod
     def extract_dataset_id(obj: dict) -> str:
         """Return unique id of the dataset."""

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

@@ -22,12 +22,13 @@ from __future__ import annotations
 from collections.abc import Sequence
 from typing import TYPE_CHECKING
 
-from airflow.exceptions import AirflowException
-from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
 from google.api_core.client_options import ClientOptions
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
 from google.cloud.aiplatform_v1 import EndpointServiceClient
 
+from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
+from airflow.providers.google.common.hooks.operation_helpers import OperationHelper
+
 if TYPE_CHECKING:
     from google.api_core.operation import Operation
     from google.api_core.retry import Retry
@@ -36,7 +37,7 @@ if TYPE_CHECKING:
     from google.protobuf.field_mask_pb2 import FieldMask
 
 
-class EndpointServiceHook(GoogleBaseHook):
+class EndpointServiceHook(GoogleBaseHook, OperationHelper):
     """Hook for Google Cloud Vertex AI Endpoint Service APIs."""
 
     def get_endpoint_service_client(self, region: str | None = None) -> EndpointServiceClient:
@@ -50,14 +51,6 @@ class EndpointServiceHook(GoogleBaseHook):
             credentials=self.get_credentials(), client_info=self.client_info, client_options=client_options
         )
 
-    def wait_for_operation(self, operation: Operation, timeout: float | None = None):
-        """Wait for long-lasting operation to complete."""
-        try:
-            return operation.result(timeout=timeout)
-        except Exception:
-            error = operation.exception(timeout=timeout)
-            raise AirflowException(error)
-
     @staticmethod
     def extract_endpoint_id(obj: dict) -> str:
         """Return unique id of the endpoint."""

airflow/providers/google/cloud/hooks/vertex_ai/feature_store.py

@@ -18,14 +18,15 @@
 
 from __future__ import annotations
 
-from airflow.exceptions import AirflowException
-from airflow.providers.google.common.consts import CLIENT_INFO
-from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID, GoogleBaseHook
 from google.api_core.client_options import ClientOptions
 from google.cloud.aiplatform_v1beta1 import (
     FeatureOnlineStoreAdminServiceClient,
 )
 
+from airflow.exceptions import AirflowException
+from airflow.providers.google.common.consts import CLIENT_INFO
+from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID, GoogleBaseHook
+
 
 class FeatureStoreHook(GoogleBaseHook):
     """

airflow/providers/google/cloud/hooks/vertex_ai/hyperparameter_tuning_job.py

@@ -29,21 +29,23 @@ import asyncio
 from collections.abc import Sequence
 from typing import TYPE_CHECKING
 
-from airflow.exceptions import AirflowException
-from airflow.providers.google.common.consts import CLIENT_INFO
-from airflow.providers.google.common.hooks.base_google import GoogleBaseAsyncHook, GoogleBaseHook
 from google.api_core.client_options import ClientOptions
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
 from google.cloud.aiplatform import CustomJob, HyperparameterTuningJob, gapic, hyperparameter_tuning
 from google.cloud.aiplatform_v1 import JobServiceAsyncClient, JobServiceClient, JobState, types
 
+from airflow.exceptions import AirflowException
+from airflow.providers.google.common.consts import CLIENT_INFO
+from airflow.providers.google.common.hooks.base_google import GoogleBaseAsyncHook, GoogleBaseHook
+from airflow.providers.google.common.hooks.operation_helpers import OperationHelper
+
 if TYPE_CHECKING:
     from google.api_core.operation import Operation
     from google.api_core.retry import AsyncRetry, Retry
     from google.cloud.aiplatform_v1.services.job_service.pagers import ListHyperparameterTuningJobsPager
 
 
-class HyperparameterTuningJobHook(GoogleBaseHook):
+class HyperparameterTuningJobHook(GoogleBaseHook, OperationHelper):
     """Hook for Google Cloud Vertex AI Hyperparameter Tuning Job APIs."""
 
     def __init__(
@@ -133,14 +135,6 @@ class HyperparameterTuningJobHook(GoogleBaseHook):
         """Return unique id of the hyperparameter_tuning_job."""
         return obj["name"].rpartition("/")[-1]
 
-    def wait_for_operation(self, operation: Operation, timeout: float | None = None):
-        """Wait for long-lasting operation to complete."""
-        try:
-            return operation.result(timeout=timeout)
-        except Exception:
-            error = operation.exception(timeout=timeout)
-            raise AirflowException(error)
-
     def cancel_hyperparameter_tuning_job(self) -> None:
         """Cancel HyperparameterTuningJob."""
         if self._hyperparameter_tuning_job:

airflow/providers/google/cloud/hooks/vertex_ai/model_service.py

@@ -23,12 +23,13 @@ from __future__ import annotations
 from collections.abc import Sequence
 from typing import TYPE_CHECKING
 
-from airflow.exceptions import AirflowException
-from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
 from google.api_core.client_options import ClientOptions
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
 from google.cloud.aiplatform_v1 import ModelServiceClient
 
+from airflow.exceptions import AirflowException
+from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
+
 if TYPE_CHECKING:
     from google.api_core.operation import Operation
     from google.api_core.retry import Retry
@@ -38,8 +39,10 @@ if TYPE_CHECKING:
     )
     from google.cloud.aiplatform_v1.types import Model, model_service
 
+from airflow.providers.google.common.hooks.operation_helpers import OperationHelper
+
 
-class ModelServiceHook(GoogleBaseHook):
+class ModelServiceHook(GoogleBaseHook, OperationHelper):
     """Hook for Google Cloud Vertex AI Endpoint Service APIs."""
 
     def get_model_service_client(self, region: str | None = None) -> ModelServiceClient:
@@ -58,14 +61,6 @@ class ModelServiceHook(GoogleBaseHook):
         """Return unique id of the model."""
         return obj["model"].rpartition("/")[-1]
 
-    def wait_for_operation(self, operation: Operation, timeout: float | None = None):
-        """Wait for long-lasting operation to complete."""
-        try:
-            return operation.result(timeout=timeout)
-        except Exception:
-            error = operation.exception(timeout=timeout)
-            raise AirflowException(error)
-
     @GoogleBaseHook.fallback_to_default_project_id
     def delete_model(
         self,

airflow/providers/google/cloud/hooks/vertex_ai/pipeline_job.py

@@ -29,9 +29,6 @@ import asyncio
 from collections.abc import Sequence
 from typing import TYPE_CHECKING, Any
 
-from airflow.exceptions import AirflowException
-from airflow.providers.google.common.consts import CLIENT_INFO
-from airflow.providers.google.common.hooks.base_google import GoogleBaseAsyncHook, GoogleBaseHook
 from google.api_core.client_options import ClientOptions
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
 from google.cloud.aiplatform import PipelineJob
@@ -42,6 +39,11 @@ from google.cloud.aiplatform_v1 import (
     types,
 )
 
+from airflow.exceptions import AirflowException
+from airflow.providers.google.common.consts import CLIENT_INFO
+from airflow.providers.google.common.hooks.base_google import GoogleBaseAsyncHook, GoogleBaseHook
+from airflow.providers.google.common.hooks.operation_helpers import OperationHelper
+
 if TYPE_CHECKING:
     from google.api_core.operation import Operation
     from google.api_core.retry import AsyncRetry, Retry
@@ -50,7 +52,7 @@ if TYPE_CHECKING:
     from google.cloud.aiplatform_v1.services.pipeline_service.pagers import ListPipelineJobsPager
 
 
-class PipelineJobHook(GoogleBaseHook):
+class PipelineJobHook(GoogleBaseHook, OperationHelper):
     """Hook for Google Cloud Vertex AI Pipeline Job APIs."""
 
     def __init__(
@@ -111,14 +113,6 @@ class PipelineJobHook(GoogleBaseHook):
             failure_policy=failure_policy,
         )
 
-    def wait_for_operation(self, operation: Operation, timeout: float | None = None):
-        """Wait for long-lasting operation to complete."""
-        try:
-            return operation.result(timeout=timeout)
-        except Exception:
-            error = operation.exception(timeout=timeout)
-            raise AirflowException(error)
-
     def cancel_pipeline_job(self) -> None:
         """Cancel PipelineJob."""
         if self._pipeline_job:

airflow/providers/google/cloud/hooks/vertex_ai/prediction_service.py

@@ -20,12 +20,13 @@ from __future__ import annotations
 from collections.abc import Sequence
 from typing import TYPE_CHECKING
 
-from airflow.providers.google.common.consts import CLIENT_INFO
-from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID, GoogleBaseHook
 from google.api_core.client_options import ClientOptions
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
 from google.cloud.aiplatform_v1 import PredictionServiceClient
 
+from airflow.providers.google.common.consts import CLIENT_INFO
+from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID, GoogleBaseHook
+
 if TYPE_CHECKING:
     from google.api_core.retry import Retry
     from google.cloud.aiplatform_v1.types import PredictResponse

airflow/providers/google/cloud/hooks/video_intelligence.py

@@ -22,8 +22,6 @@ from __future__ import annotations
 from collections.abc import Sequence
 from typing import TYPE_CHECKING
 
-from airflow.providers.google.common.consts import CLIENT_INFO
-from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
 from google.cloud.videointelligence_v1 import (
     Feature,
@@ -31,6 +29,9 @@ from google.cloud.videointelligence_v1 import (
     VideoIntelligenceServiceClient,
 )
 
+from airflow.providers.google.common.consts import CLIENT_INFO
+from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
+
 if TYPE_CHECKING:
     from google.api_core.operation import Operation
     from google.api_core.retry import Retry

airflow/providers/google/cloud/hooks/vision.py

@@ -24,9 +24,6 @@ from copy import deepcopy
 from functools import cached_property
 from typing import TYPE_CHECKING, Any, Callable
 
-from airflow.exceptions import AirflowException
-from airflow.providers.google.common.consts import CLIENT_INFO
-from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID, GoogleBaseHook
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
 from google.cloud.vision_v1 import (
     AnnotateImageRequest,
@@ -39,6 +36,10 @@ from google.cloud.vision_v1 import (
 )
 from google.protobuf.json_format import MessageToDict
 
+from airflow.exceptions import AirflowException
+from airflow.providers.google.common.consts import CLIENT_INFO
+from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID, GoogleBaseHook
+
 if TYPE_CHECKING:
     from google.api_core.retry import Retry
     from google.protobuf import field_mask_pb2

airflow/providers/google/cloud/hooks/workflows.py

@@ -19,12 +19,13 @@ from __future__ import annotations
 from collections.abc import Sequence
 from typing import TYPE_CHECKING
 
-from airflow.providers.google.common.consts import CLIENT_INFO
-from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID, GoogleBaseHook
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
 from google.cloud.workflows.executions_v1beta import Execution, ExecutionsClient
 from google.cloud.workflows_v1beta import Workflow, WorkflowsClient
 
+from airflow.providers.google.common.consts import CLIENT_INFO
+from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID, GoogleBaseHook
+
 if TYPE_CHECKING:
     from google.api_core.operation import Operation
     from google.api_core.retry import Retry