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

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__(
@@ -61,17 +63,9 @@ class BatchPredictionJobHook(GoogleBaseHook):
             client_options = ClientOptions()
 
         return JobServiceClient(
-            credentials=self.get_credentials(), client_info=self.client_info, client_options=client_options
+            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_batch_prediction_job_id(obj: dict) -> str:
         """Return unique id of the batch_prediction_job."""
@@ -116,7 +110,7 @@ class BatchPredictionJobHook(GoogleBaseHook):
         :param project_id: Required. Project to run training in.
         :param region: Required. Location to run training in.
         :param job_display_name: Required. The user-defined name of the BatchPredictionJob. The name can be
-            up to 128 characters long and can be consist of any UTF-8 characters.
+            up to 128 characters long and can consist of any UTF-8 characters.
         :param model_name: Required. A fully-qualified model resource name or model ID.
         :param instances_format: Required. The format in which instances are provided. Must be one of the
             formats listed in `Model.supported_input_storage_formats`. Default is "jsonl" when using
@@ -273,7 +267,7 @@ class BatchPredictionJobHook(GoogleBaseHook):
         :param project_id: Required. Project to run training in.
         :param region: Required. Location to run training in.
         :param job_display_name: Required. The user-defined name of the BatchPredictionJob. The name can be
-            up to 128 characters long and can be consist of any UTF-8 characters.
+            up to 128 characters long and can consist of any UTF-8 characters.
         :param model_name: Required. A fully-qualified model resource name or model ID.
         :param instances_format: Required. The format in which instances are provided. Must be one of the
             formats listed in `Model.supported_input_storage_formats`. Default is "jsonl" when using

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,19 +42,23 @@ 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
     from google.auth.credentials import Credentials
     from google.cloud.aiplatform_v1.services.job_service.pagers import ListCustomJobsPager
     from google.cloud.aiplatform_v1.services.pipeline_service.pagers import (
-        ListPipelineJobsPager,
         ListTrainingPipelinesPager,
     )
-    from google.cloud.aiplatform_v1.types import CustomJob, PipelineJob, TrainingPipeline
+    from google.cloud.aiplatform_v1.types import CustomJob, PscInterfaceConfig, TrainingPipeline
 
 
-class CustomJobHook(GoogleBaseHook):
+class CustomJobHook(GoogleBaseHook, OperationHelper):
     """Hook for Google Cloud Vertex AI Custom Job APIs."""
 
     def __init__(
@@ -276,14 +276,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:
@@ -325,6 +317,7 @@ class CustomJobHook(GoogleBaseHook):
         is_default_version: bool | None = None,
         model_version_aliases: list[str] | None = None,
         model_version_description: str | None = None,
+        psc_interface_config: PscInterfaceConfig | None = None,
     ) -> tuple[models.Model | None, str, str]:
         """Run a training pipeline job and wait until its completion."""
         model = job.run(
@@ -358,6 +351,7 @@ class CustomJobHook(GoogleBaseHook):
             is_default_version=is_default_version,
             model_version_aliases=model_version_aliases,
             model_version_description=model_version_description,
+            psc_interface_config=psc_interface_config,
         )
         training_id = self.extract_training_id(job.resource_name)
         custom_job_id = self.extract_custom_job_id(
@@ -374,54 +368,6 @@ class CustomJobHook(GoogleBaseHook):
             )
         return model, training_id, custom_job_id
 
-    @GoogleBaseHook.fallback_to_default_project_id
-    @deprecated(
-        planned_removal_date="March 01, 2025",
-        use_instead="PipelineJobHook.cancel_pipeline_job",
-        category=AirflowProviderDeprecationWarning,
-    )
-    def cancel_pipeline_job(
-        self,
-        project_id: str,
-        region: str,
-        pipeline_job: str,
-        retry: Retry | _MethodDefault = DEFAULT,
-        timeout: float | None = None,
-        metadata: Sequence[tuple[str, str]] = (),
-    ) -> None:
-        """
-        Cancel a PipelineJob.
-
-        Starts asynchronous cancellation on the PipelineJob. The server makes the best
-        effort to cancel the pipeline, but success is not guaranteed. Clients can use
-        [PipelineService.GetPipelineJob][google.cloud.aiplatform.v1.PipelineService.GetPipelineJob] or other
-        methods to check whether the cancellation succeeded or whether the pipeline completed despite
-        cancellation. On successful cancellation, the PipelineJob is not deleted; instead it becomes a
-        pipeline with a [PipelineJob.error][google.cloud.aiplatform.v1.PipelineJob.error] value with a
-        [google.rpc.Status.code][google.rpc.Status.code] of 1, corresponding to ``Code.CANCELLED``, and
-        [PipelineJob.state][google.cloud.aiplatform.v1.PipelineJob.state] is set to ``CANCELLED``.
-
-        This method is deprecated, please use `PipelineJobHook.cancel_pipeline_job` method.
-
-        :param project_id: Required. The ID of the Google Cloud project that the service belongs to.
-        :param region: Required. The ID of the Google Cloud region that the service belongs to.
-        :param pipeline_job: The name of the PipelineJob to cancel.
-        :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.
-        """
-        client = self.get_pipeline_service_client(region)
-        name = client.pipeline_job_path(project_id, region, pipeline_job)
-
-        client.cancel_pipeline_job(
-            request={
-                "name": name,
-            },
-            retry=retry,
-            timeout=timeout,
-            metadata=metadata,
-        )
-
     @GoogleBaseHook.fallback_to_default_project_id
     def cancel_training_pipeline(
         self,
@@ -504,53 +450,6 @@ class CustomJobHook(GoogleBaseHook):
             metadata=metadata,
         )
 
-    @GoogleBaseHook.fallback_to_default_project_id
-    @deprecated(
-        planned_removal_date="March 01, 2025",
-        use_instead="PipelineJobHook.create_pipeline_job",
-        category=AirflowProviderDeprecationWarning,
-    )
-    def create_pipeline_job(
-        self,
-        project_id: str,
-        region: str,
-        pipeline_job: PipelineJob,
-        pipeline_job_id: str,
-        retry: Retry | _MethodDefault = DEFAULT,
-        timeout: float | None = None,
-        metadata: Sequence[tuple[str, str]] = (),
-    ) -> PipelineJob:
-        """
-        Create a PipelineJob. A PipelineJob will run immediately when created.
-
-        This method is deprecated, please use `PipelineJobHook.create_pipeline_job` method.
-
-        :param project_id: Required. The ID of the Google Cloud project that the service belongs to.
-        :param region: Required. The ID of the Google Cloud region that the service belongs to.
-        :param pipeline_job:  Required. The PipelineJob to create.
-        :param pipeline_job_id:  The ID to use for the PipelineJob, which will become the final component of
-            the PipelineJob name. If not provided, an ID will be automatically generated.
-
-            This value should be less than 128 characters, and valid characters are /[a-z][0-9]-/.
-        :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.
-        """
-        client = self.get_pipeline_service_client(region)
-        parent = client.common_location_path(project_id, region)
-
-        result = client.create_pipeline_job(
-            request={
-                "parent": parent,
-                "pipeline_job": pipeline_job,
-                "pipeline_job_id": pipeline_job_id,
-            },
-            retry=retry,
-            timeout=timeout,
-            metadata=metadata,
-        )
-        return result
-
     @GoogleBaseHook.fallback_to_default_project_id
     def create_training_pipeline(
         self,
@@ -677,6 +576,7 @@ class CustomJobHook(GoogleBaseHook):
         timestamp_split_column_name: str | None = None,
         tensorboard: str | None = None,
         sync=True,
+        psc_interface_config: PscInterfaceConfig | None = None,
     ) -> tuple[models.Model | None, str, str]:
         """
         Create Custom Container Training Job.
@@ -940,6 +840,8 @@ class CustomJobHook(GoogleBaseHook):
         :param sync: Whether to execute the AI Platform job synchronously. If False, this method
                 will be executed in concurrent Future and any downstream object will
                 be immediately returned and synced when the Future has completed.
+        :param psc_interface_config: Optional. Configuration for Private Service Connect interface used for
+            training.
         """
         self._job = self.get_custom_container_training_job(
             project=project_id,
@@ -999,6 +901,7 @@ class CustomJobHook(GoogleBaseHook):
             is_default_version=is_default_version,
             model_version_aliases=model_version_aliases,
             model_version_description=model_version_description,
+            psc_interface_config=psc_interface_config,
         )
 
         return model, training_id, custom_job_id
@@ -1061,6 +964,7 @@ class CustomJobHook(GoogleBaseHook):
         model_version_aliases: list[str] | None = None,
         model_version_description: str | None = None,
         sync=True,
+        psc_interface_config: PscInterfaceConfig | None = None,
     ) -> tuple[models.Model | None, str, str]:
         """
         Create Custom Python Package Training Job.
@@ -1323,6 +1227,8 @@ class CustomJobHook(GoogleBaseHook):
         :param sync: Whether to execute the AI Platform job synchronously. If False, this method
                 will be executed in concurrent Future and any downstream object will
                 be immediately returned and synced when the Future has completed.
+        :param psc_interface_config: Optional. Configuration for Private Service Connect interface used for
+            training.
         """
         self._job = self.get_custom_python_package_training_job(
             project=project_id,
@@ -1383,6 +1289,7 @@ class CustomJobHook(GoogleBaseHook):
             is_default_version=is_default_version,
             model_version_aliases=model_version_aliases,
             model_version_description=model_version_description,
+            psc_interface_config=psc_interface_config,
         )
 
         return model, training_id, custom_job_id
@@ -1445,6 +1352,7 @@ class CustomJobHook(GoogleBaseHook):
         timestamp_split_column_name: str | None = None,
         tensorboard: str | None = None,
         sync=True,
+        psc_interface_config: PscInterfaceConfig | None = None,
     ) -> tuple[models.Model | None, str, str]:
         """
         Create Custom Training Job.
@@ -1707,6 +1615,8 @@ class CustomJobHook(GoogleBaseHook):
         :param sync: Whether to execute the AI Platform job synchronously. If False, this method
                 will be executed in concurrent Future and any downstream object will
                 be immediately returned and synced when the Future has completed.
+        :param psc_interface_config: Optional. Configuration for Private Service Connect interface used for
+            training.
         """
         self._job = self.get_custom_training_job(
             project=project_id,
@@ -1767,6 +1677,7 @@ class CustomJobHook(GoogleBaseHook):
             is_default_version=is_default_version,
             model_version_aliases=model_version_aliases,
             model_version_description=model_version_description,
+            psc_interface_config=psc_interface_config,
         )
 
         return model, training_id, custom_job_id
@@ -1828,6 +1739,7 @@ class CustomJobHook(GoogleBaseHook):
         predefined_split_column_name: str | None = None,
         timestamp_split_column_name: str | None = None,
         tensorboard: str | None = None,
+        psc_interface_config: PscInterfaceConfig | None = None,
     ) -> CustomContainerTrainingJob:
         """
         Create and submit a Custom Container Training Job pipeline, then exit without waiting for it to complete.
@@ -2088,6 +2000,8 @@ class CustomJobHook(GoogleBaseHook):
                 ``projects/{project}/locations/{location}/tensorboards/{tensorboard}``
                 For more information on configuring your service account please visit:
                 https://cloud.google.com/vertex-ai/docs/experiments/tensorboard-training
+        :param psc_interface_config: Optional. Configuration for Private Service Connect interface used for
+            training.
         """
         self._job = self.get_custom_container_training_job(
             project=project_id,
@@ -2146,6 +2060,7 @@ class CustomJobHook(GoogleBaseHook):
             model_version_aliases=model_version_aliases,
             model_version_description=model_version_description,
             sync=False,
+            psc_interface_config=psc_interface_config,
         )
         return self._job
 
@@ -2207,6 +2122,7 @@ class CustomJobHook(GoogleBaseHook):
         is_default_version: bool | None = None,
         model_version_aliases: list[str] | None = None,
         model_version_description: str | None = None,
+        psc_interface_config: PscInterfaceConfig | None = None,
     ) -> CustomPythonPackageTrainingJob:
         """
         Create and submit a Custom Python Package Training Job pipeline, then exit without waiting for it to complete.
@@ -2466,6 +2382,8 @@ class CustomJobHook(GoogleBaseHook):
                 ``projects/{project}/locations/{location}/tensorboards/{tensorboard}``
                 For more information on configuring your service account please visit:
                 https://cloud.google.com/vertex-ai/docs/experiments/tensorboard-training
+        :param psc_interface_config: Optional. Configuration for Private Service Connect interface used for
+            training.
         """
         self._job = self.get_custom_python_package_training_job(
             project=project_id,
@@ -2525,6 +2443,7 @@ class CustomJobHook(GoogleBaseHook):
             model_version_aliases=model_version_aliases,
             model_version_description=model_version_description,
             sync=False,
+            psc_interface_config=psc_interface_config,
         )
 
         return self._job
@@ -2587,6 +2506,7 @@ class CustomJobHook(GoogleBaseHook):
         predefined_split_column_name: str | None = None,
         timestamp_split_column_name: str | None = None,
         tensorboard: str | None = None,
+        psc_interface_config: PscInterfaceConfig | None = None,
     ) -> CustomTrainingJob:
         """
         Create and submit a Custom Training Job pipeline, then exit without waiting for it to complete.
@@ -2850,6 +2770,8 @@ class CustomJobHook(GoogleBaseHook):
                 ``projects/{project}/locations/{location}/tensorboards/{tensorboard}``
                 For more information on configuring your service account please visit:
                 https://cloud.google.com/vertex-ai/docs/experiments/tensorboard-training
+        :param psc_interface_config: Optional. Configuration for Private Service Connect interface used for
+            training.
         """
         self._job = self.get_custom_training_job(
             project=project_id,
@@ -2909,6 +2831,7 @@ class CustomJobHook(GoogleBaseHook):
             model_version_aliases=model_version_aliases,
             model_version_description=model_version_description,
             sync=False,
+            psc_interface_config=psc_interface_config,
         )
         return self._job
 
@@ -2976,46 +2899,6 @@ class CustomJobHook(GoogleBaseHook):
         )
         return result
 
-    @GoogleBaseHook.fallback_to_default_project_id
-    @deprecated(
-        planned_removal_date="March 01, 2025",
-        use_instead="PipelineJobHook.get_pipeline_job",
-        category=AirflowProviderDeprecationWarning,
-    )
-    def get_pipeline_job(
-        self,
-        project_id: str,
-        region: str,
-        pipeline_job: str,
-        retry: Retry | _MethodDefault = DEFAULT,
-        timeout: float | None = None,
-        metadata: Sequence[tuple[str, str]] = (),
-    ) -> PipelineJob:
-        """
-        Get a PipelineJob.
-
-        This method is deprecated, please use `PipelineJobHook.get_pipeline_job` method.
-
-        :param project_id: Required. The ID of the Google Cloud project that the service belongs to.
-        :param region: Required. The ID of the Google Cloud region that the service belongs to.
-        :param pipeline_job: Required. The name of the PipelineJob resource.
-        :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.
-        """
-        client = self.get_pipeline_service_client(region)
-        name = client.pipeline_job_path(project_id, region, pipeline_job)
-
-        result = client.get_pipeline_job(
-            request={
-                "name": name,
-            },
-            retry=retry,
-            timeout=timeout,
-            metadata=metadata,
-        )
-        return result
-
     @GoogleBaseHook.fallback_to_default_project_id
     def get_training_pipeline(
         self,
@@ -3082,101 +2965,6 @@ class CustomJobHook(GoogleBaseHook):
         )
         return result
 
-    @GoogleBaseHook.fallback_to_default_project_id
-    @deprecated(
-        planned_removal_date="March 01, 2025",
-        use_instead="PipelineJobHook.list_pipeline_jobs",
-        category=AirflowProviderDeprecationWarning,
-    )
-    def list_pipeline_jobs(
-        self,
-        project_id: str,
-        region: str,
-        page_size: int | None = None,
-        page_token: str | None = None,
-        filter: str | None = None,
-        order_by: str | None = None,
-        retry: Retry | _MethodDefault = DEFAULT,
-        timeout: float | None = None,
-        metadata: Sequence[tuple[str, str]] = (),
-    ) -> ListPipelineJobsPager:
-        """
-        List PipelineJobs in a Location.
-
-        This method is deprecated, please use `PipelineJobHook.list_pipeline_jobs` method.
-
-        :param project_id: Required. The ID of the Google Cloud project that the service belongs to.
-        :param region: Required. The ID of the Google Cloud region that the service belongs to.
-        :param filter: Optional. Lists the PipelineJobs that match the filter expression. The
-            following fields are supported:
-
-            -  ``pipeline_name``: Supports ``=`` and ``!=`` comparisons.
-            -  ``display_name``: Supports ``=``, ``!=`` comparisons, and
-               ``:`` wildcard.
-            -  ``pipeline_job_user_id``: Supports ``=``, ``!=``
-               comparisons, and ``:`` wildcard. for example, can check
-               if pipeline's display_name contains *step* by doing
-               display_name:"*step*"
-            -  ``create_time``: Supports ``=``, ``!=``, ``<``, ``>``,
-               ``<=``, and ``>=`` comparisons. Values must be in RFC
-               3339 format.
-            -  ``update_time``: Supports ``=``, ``!=``, ``<``, ``>``,
-               ``<=``, and ``>=`` comparisons. Values must be in RFC
-               3339 format.
-            -  ``end_time``: Supports ``=``, ``!=``, ``<``, ``>``,
-               ``<=``, and ``>=`` comparisons. Values must be in RFC
-               3339 format.
-            -  ``labels``: Supports key-value equality and key presence.
-
-            Filter expressions can be combined together using logical
-            operators (``AND`` & ``OR``). For example:
-            ``pipeline_name="test" AND create_time>"2020-05-18T13:30:00Z"``.
-
-            The syntax to define filter expression is based on
-            https://google.aip.dev/160.
-        :param page_size: Optional. The standard list page size.
-        :param page_token: Optional. The standard list page token. Typically obtained via
-            [ListPipelineJobsResponse.next_page_token][google.cloud.aiplatform.v1.ListPipelineJobsResponse.next_page_token]
-            of the previous
-            [PipelineService.ListPipelineJobs][google.cloud.aiplatform.v1.PipelineService.ListPipelineJobs]
-            call.
-        :param order_by: Optional. A comma-separated list of fields to order by. The default
-            sort order is in ascending order. Use "desc" after a field
-            name for descending. You can have multiple order_by fields
-            provided e.g. "create_time desc, end_time", "end_time,
-            start_time, update_time" For example, using "create_time
-            desc, end_time" will order results by create time in
-            descending order, and if there are multiple jobs having the
-            same create time, order them by the end time in ascending
-            order. if order_by is not specified, it will order by
-            default order is create time in descending order. Supported
-            fields:
-
-            -  ``create_time``
-            -  ``update_time``
-            -  ``end_time``
-            -  ``start_time``
-        :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.
-        """
-        client = self.get_pipeline_service_client(region)
-        parent = client.common_location_path(project_id, region)
-
-        result = client.list_pipeline_jobs(
-            request={
-                "parent": parent,
-                "page_size": page_size,
-                "page_token": page_token,
-                "filter": filter,
-                "order_by": order_by,
-            },
-            retry=retry,
-            timeout=timeout,
-            metadata=metadata,
-        )
-        return result
-
     @GoogleBaseHook.fallback_to_default_project_id
     def list_training_pipelines(
         self,
@@ -3299,44 +3087,6 @@ class CustomJobHook(GoogleBaseHook):
         )
         return result
 
-    @GoogleBaseHook.fallback_to_default_project_id
-    @deprecated(
-        planned_removal_date="March 01, 2025",
-        use_instead="PipelineJobHook.delete_pipeline_job",
-        category=AirflowProviderDeprecationWarning,
-    )
-    def delete_pipeline_job(
-        self,
-        project_id: str,
-        region: str,
-        pipeline_job: str,
-        retry: Retry | _MethodDefault = DEFAULT,
-        timeout: float | None = None,
-        metadata: Sequence[tuple[str, str]] = (),
-    ) -> Operation:
-        """
-        Delete a PipelineJob.
-
-        This method is deprecated, please use `PipelineJobHook.delete_pipeline_job` method.
-
-        :param project_id: Required. The ID of the Google Cloud project that the service belongs to.
-        :param region: Required. The ID of the Google Cloud region that the service belongs to.
-        :param pipeline_job: Required. The name of the PipelineJob resource to be deleted.
-        :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.
-        """
-        client = self.get_pipeline_service_client(region)
-        name = client.pipeline_job_path(project_id, region, pipeline_job)
-
-        result = client.delete_pipeline_job(
-            request={"name": name},
-            retry=retry,
-            timeout=timeout,
-            metadata=metadata,
-        )
-        return result
-
 
 class CustomJobAsyncHook(GoogleBaseAsyncHook):
     """Async hook for Custom Job Service Client."""

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,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.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.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
@@ -36,7 +38,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:
@@ -47,17 +49,9 @@ class EndpointServiceHook(GoogleBaseHook):
             client_options = ClientOptions()
 
         return EndpointServiceClient(
-            credentials=self.get_credentials(), client_info=self.client_info, client_options=client_options
+            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_endpoint_id(obj: dict) -> str:
         """Return unique id of the endpoint."""