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

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

@@ -25,10 +25,12 @@ import time
 from typing import Any, AsyncIterator, Sequence
 
 from google.api_core.exceptions import NotFound
-from google.cloud.dataproc_v1 import Batch, ClusterStatus, JobStatus
+from google.cloud.dataproc_v1 import Batch, Cluster, ClusterStatus, JobStatus
 
-from airflow.providers.google.cloud.hooks.dataproc import DataprocAsyncHook
+from airflow.exceptions import AirflowException
+from airflow.providers.google.cloud.hooks.dataproc import DataprocAsyncHook, DataprocHook
 from airflow.providers.google.cloud.utils.dataproc import DataprocOperationType
+from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
 from airflow.triggers.base import BaseTrigger, TriggerEvent
 
 
@@ -38,10 +40,12 @@ class DataprocBaseTrigger(BaseTrigger):
     def __init__(
         self,
         region: str,
-        project_id: str | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
         gcp_conn_id: str = "google_cloud_default",
         impersonation_chain: str | Sequence[str] | None = None,
         polling_interval_seconds: int = 30,
+        cancel_on_kill: bool = True,
+        delete_on_error: bool = True,
     ):
         super().__init__()
         self.region = region
@@ -49,6 +53,8 @@ class DataprocBaseTrigger(BaseTrigger):
         self.gcp_conn_id = gcp_conn_id
         self.impersonation_chain = impersonation_chain
         self.polling_interval_seconds = polling_interval_seconds
+        self.cancel_on_kill = cancel_on_kill
+        self.delete_on_error = delete_on_error
 
     def get_async_hook(self):
         return DataprocAsyncHook(
@@ -56,6 +62,16 @@ class DataprocBaseTrigger(BaseTrigger):
             impersonation_chain=self.impersonation_chain,
         )
 
+    def get_sync_hook(self):
+        # The synchronous hook is utilized to delete the cluster when a task is cancelled.
+        # This is because the asynchronous hook deletion is not awaited when the trigger task
+        # is cancelled. The call for deleting the cluster or job through the sync hook is not a blocking
+        # call, which means it does not wait until the cluster or job is deleted.
+        return DataprocHook(
+            gcp_conn_id=self.gcp_conn_id,
+            impersonation_chain=self.impersonation_chain,
+        )
+
 
 class DataprocSubmitTrigger(DataprocBaseTrigger):
     """
@@ -90,20 +106,39 @@ class DataprocSubmitTrigger(DataprocBaseTrigger):
                 "gcp_conn_id": self.gcp_conn_id,
                 "impersonation_chain": self.impersonation_chain,
                 "polling_interval_seconds": self.polling_interval_seconds,
+                "cancel_on_kill": self.cancel_on_kill,
             },
         )
 
     async def run(self):
-        while True:
-            job = await self.get_async_hook().get_job(
-                project_id=self.project_id, region=self.region, job_id=self.job_id
-            )
-            state = job.status.state
-            self.log.info("Dataproc job: %s is in state: %s", self.job_id, state)
-            if state in (JobStatus.State.DONE, JobStatus.State.CANCELLED, JobStatus.State.ERROR):
-                break
-            await asyncio.sleep(self.polling_interval_seconds)
-        yield TriggerEvent({"job_id": self.job_id, "job_state": state, "job": job})
+        try:
+            while True:
+                job = await self.get_async_hook().get_job(
+                    project_id=self.project_id, region=self.region, job_id=self.job_id
+                )
+                state = job.status.state
+                self.log.info("Dataproc job: %s is in state: %s", self.job_id, state)
+                if state in (JobStatus.State.DONE, JobStatus.State.CANCELLED, JobStatus.State.ERROR):
+                    break
+                await asyncio.sleep(self.polling_interval_seconds)
+            yield TriggerEvent({"job_id": self.job_id, "job_state": state, "job": job})
+        except asyncio.CancelledError:
+            self.log.info("Task got cancelled.")
+            try:
+                if self.job_id and self.cancel_on_kill:
+                    self.log.info("Cancelling the job: %s", self.job_id)
+                    # The synchronous hook is utilized to delete the cluster when a task is cancelled. This
+                    # is because the asynchronous hook deletion is not awaited when the trigger task is
+                    # cancelled. The call for deleting the cluster or job through the sync hook is not a
+                    # blocking call, which means it does not wait until the cluster or job is deleted.
+                    self.get_sync_hook().cancel_job(
+                        job_id=self.job_id, project_id=self.project_id, region=self.region
+                    )
+                    self.log.info("Job: %s is cancelled", self.job_id)
+                    yield TriggerEvent({"job_id": self.job_id, "job_state": ClusterStatus.State.DELETING})
+            except Exception as e:
+                self.log.error("Failed to cancel the job: %s with error : %s", self.job_id, str(e))
+                raise e
 
 
 class DataprocClusterTrigger(DataprocBaseTrigger):
@@ -139,24 +174,73 @@ class DataprocClusterTrigger(DataprocBaseTrigger):
                 "gcp_conn_id": self.gcp_conn_id,
                 "impersonation_chain": self.impersonation_chain,
                 "polling_interval_seconds": self.polling_interval_seconds,
+                "delete_on_error": self.delete_on_error,
             },
         )
 
     async def run(self) -> AsyncIterator[TriggerEvent]:
-        while True:
-            cluster = await self.get_async_hook().get_cluster(
-                project_id=self.project_id, region=self.region, cluster_name=self.cluster_name
+        try:
+            while True:
+                cluster = await self.fetch_cluster()
+                state = cluster.status.state
+                if state == ClusterStatus.State.ERROR:
+                    await self.delete_when_error_occurred(cluster)
+                    yield TriggerEvent(
+                        {
+                            "cluster_name": self.cluster_name,
+                            "cluster_state": ClusterStatus.State.DELETING,
+                            "cluster": cluster,
+                        }
+                    )
+                    return
+                elif state == ClusterStatus.State.RUNNING:
+                    yield TriggerEvent(
+                        {
+                            "cluster_name": self.cluster_name,
+                            "cluster_state": state,
+                            "cluster": cluster,
+                        }
+                    )
+                    return
+                self.log.info("Current state is %s", state)
+                self.log.info("Sleeping for %s seconds.", self.polling_interval_seconds)
+                await asyncio.sleep(self.polling_interval_seconds)
+        except asyncio.CancelledError:
+            try:
+                if self.delete_on_error:
+                    self.log.info("Deleting cluster %s.", self.cluster_name)
+                    # The synchronous hook is utilized to delete the cluster when a task is cancelled.
+                    # This is because the asynchronous hook deletion is not awaited when the trigger task
+                    # is cancelled. The call for deleting the cluster through the sync hook is not a blocking
+                    # call, which means it does not wait until the cluster is deleted.
+                    self.get_sync_hook().delete_cluster(
+                        region=self.region, cluster_name=self.cluster_name, project_id=self.project_id
+                    )
+                    self.log.info("Deleted cluster %s during cancellation.", self.cluster_name)
+            except Exception as e:
+                self.log.error("Error during cancellation handling: %s", e)
+                raise AirflowException("Error during cancellation handling: %s", e)
+
+    async def fetch_cluster(self) -> Cluster:
+        """Fetch the cluster status."""
+        return await self.get_async_hook().get_cluster(
+            project_id=self.project_id, region=self.region, cluster_name=self.cluster_name
+        )
+
+    async def delete_when_error_occurred(self, cluster: Cluster) -> None:
+        """
+        Delete the cluster on error.
+
+        :param cluster: The cluster to delete.
+        """
+        if self.delete_on_error:
+            self.log.info("Deleting cluster %s.", self.cluster_name)
+            await self.get_async_hook().delete_cluster(
+                region=self.region, cluster_name=self.cluster_name, project_id=self.project_id
             )
-            state = cluster.status.state
-            self.log.info("Dataproc cluster: %s is in state: %s", self.cluster_name, state)
-            if state in (
-                ClusterStatus.State.ERROR,
-                ClusterStatus.State.RUNNING,
-            ):
-                break
-            self.log.info("Sleeping for %s seconds.", self.polling_interval_seconds)
-            await asyncio.sleep(self.polling_interval_seconds)
-        yield TriggerEvent({"cluster_name": self.cluster_name, "cluster_state": state, "cluster": cluster})
+            self.log.info("Cluster %s has been deleted.", self.cluster_name)
+        else:
+            self.log.info("Cluster %s is not deleted as delete_on_error is set to False.", self.cluster_name)
 
 
 class DataprocBatchTrigger(DataprocBaseTrigger):

airflow/providers/google/cloud/triggers/mlengine.py

@@ -20,6 +20,7 @@ import asyncio
 from typing import Any, AsyncIterator, Sequence
 
 from airflow.providers.google.cloud.hooks.mlengine import MLEngineAsyncHook
+from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
 from airflow.triggers.base import BaseTrigger, TriggerEvent
 
 
@@ -45,7 +46,7 @@ class MLEngineStartTrainingJobTrigger(BaseTrigger):
         runtime_version: str | None = None,
         python_version: str | None = None,
         job_dir: str | None = None,
-        project_id: str | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
         labels: dict[str, str] | None = None,
         gcp_conn_id: str = "google_cloud_default",
         impersonation_chain: str | Sequence[str] | None = None,

airflow/providers/google/cloud/triggers/vertex_ai.py

@@ -29,6 +29,7 @@ from google.cloud.aiplatform_v1 import (
 
 from airflow.exceptions import AirflowException
 from airflow.providers.google.cloud.hooks.vertex_ai.batch_prediction_job import BatchPredictionJobAsyncHook
+from airflow.providers.google.cloud.hooks.vertex_ai.custom_job import CustomJobAsyncHook
 from airflow.providers.google.cloud.hooks.vertex_ai.hyperparameter_tuning_job import (
     HyperparameterTuningJobAsyncHook,
 )
@@ -189,3 +190,96 @@ class RunPipelineJobTrigger(BaseVertexAIJobTrigger):
             poll_interval=self.poll_interval,
         )
         return job
+
+
+class CustomTrainingJobTrigger(BaseVertexAIJobTrigger):
+    """
+    Make async calls to Vertex AI to check the state of a running custom training job.
+
+    Return the job when it enters a completed state.
+    """
+
+    job_type_verbose_name = "Custom Training Job"
+    job_serializer_class = types.TrainingPipeline
+    statuses_success = {
+        PipelineState.PIPELINE_STATE_PAUSED,
+        PipelineState.PIPELINE_STATE_SUCCEEDED,
+    }
+
+    @cached_property
+    def async_hook(self) -> CustomJobAsyncHook:
+        return CustomJobAsyncHook(
+            gcp_conn_id=self.conn_id,
+            impersonation_chain=self.impersonation_chain,
+        )
+
+    async def _wait_job(self) -> types.TrainingPipeline:
+        pipeline: types.TrainingPipeline = await self.async_hook.wait_for_training_pipeline(
+            project_id=self.project_id,
+            location=self.location,
+            pipeline_id=self.job_id,
+            poll_interval=self.poll_interval,
+        )
+        return pipeline
+
+
+class CustomContainerTrainingJobTrigger(BaseVertexAIJobTrigger):
+    """
+    Make async calls to Vertex AI to check the state of a running custom container training job.
+
+    Return the job when it enters a completed state.
+    """
+
+    job_type_verbose_name = "Custom Container Training Job"
+    job_serializer_class = types.TrainingPipeline
+    statuses_success = {
+        PipelineState.PIPELINE_STATE_PAUSED,
+        PipelineState.PIPELINE_STATE_SUCCEEDED,
+    }
+
+    @cached_property
+    def async_hook(self) -> CustomJobAsyncHook:
+        return CustomJobAsyncHook(
+            gcp_conn_id=self.conn_id,
+            impersonation_chain=self.impersonation_chain,
+        )
+
+    async def _wait_job(self) -> types.TrainingPipeline:
+        pipeline: types.TrainingPipeline = await self.async_hook.wait_for_training_pipeline(
+            project_id=self.project_id,
+            location=self.location,
+            pipeline_id=self.job_id,
+            poll_interval=self.poll_interval,
+        )
+        return pipeline
+
+
+class CustomPythonPackageTrainingJobTrigger(BaseVertexAIJobTrigger):
+    """
+    Make async calls to Vertex AI to check the state of a running custom python package training job.
+
+    Return the job when it enters a completed state.
+    """
+
+    job_type_verbose_name = "Custom Python Package Training Job"
+    job_serializer_class = types.TrainingPipeline
+    statuses_success = {
+        PipelineState.PIPELINE_STATE_PAUSED,
+        PipelineState.PIPELINE_STATE_SUCCEEDED,
+    }
+
+    @cached_property
+    def async_hook(self) -> CustomJobAsyncHook:
+        return CustomJobAsyncHook(
+            gcp_conn_id=self.conn_id,
+            impersonation_chain=self.impersonation_chain,
+        )
+
+    async def _wait_job(self) -> types.TrainingPipeline:
+        pipeline: types.TrainingPipeline = await self.async_hook.wait_for_training_pipeline(
+            project_id=self.project_id,
+            location=self.location,
+            pipeline_id=self.job_id,
+            poll_interval=self.poll_interval,
+        )
+        return pipeline

airflow/providers/google/common/hooks/base_google.py

@@ -114,6 +114,19 @@ def is_operation_in_progress_exception(exception: Exception) -> bool:
     return False
 
 
+def is_refresh_credentials_exception(exception: Exception) -> bool:
+    """
+    Handle refresh credentials exceptions.
+
+    Some calls return 502 (server error) in case a new token cannot be obtained.
+
+    * Google BigQuery
+    """
+    if isinstance(exception, RefreshError):
+        return "Unable to acquire impersonated credentials" in str(exception)
+    return False
+
+
 class retry_if_temporary_quota(tenacity.retry_if_exception):
     """Retries if there was an exception for exceeding the temporary quote limit."""
 
@@ -122,12 +135,19 @@ class retry_if_temporary_quota(tenacity.retry_if_exception):
 
 
 class retry_if_operation_in_progress(tenacity.retry_if_exception):
-    """Retries if there was an exception for exceeding the temporary quote limit."""
+    """Retries if there was an exception in case of operation in progress."""
 
     def __init__(self):
         super().__init__(is_operation_in_progress_exception)
 
 
+class retry_if_temporary_refresh_credentials(tenacity.retry_if_exception):
+    """Retries if there was an exception for refreshing credentials."""
+
+    def __init__(self):
+        super().__init__(is_refresh_credentials_exception)
+
+
 # A fake project_id to use in functions decorated by fallback_to_default_project_id
 # This allows the 'project_id' argument to be of type str instead of str | None,
 # making it easier to type hint the function body without dealing with the None
@@ -364,14 +384,14 @@ class GoogleBaseHook(BaseHook):
         return hasattr(self, "extras") and get_field(self.extras, f) or default
 
     @property
-    def project_id(self) -> str | None:
+    def project_id(self) -> str:
         """
         Returns project id.
 
         :return: id of the project
         """
         _, project_id = self.get_credentials_and_project_id()
-        return project_id
+        return project_id or PROVIDE_PROJECT_ID
 
     @property
     def num_retries(self) -> int:
@@ -426,7 +446,7 @@ class GoogleBaseHook(BaseHook):
     def quota_retry(*args, **kwargs) -> Callable:
         """Provide a mechanism to repeat requests in response to exceeding a temporary quota limit."""
 
-        def decorator(fun: Callable):
+        def decorator(func: Callable):
             default_kwargs = {
                 "wait": tenacity.wait_exponential(multiplier=1, max=100),
                 "retry": retry_if_temporary_quota(),
@@ -434,7 +454,7 @@ class GoogleBaseHook(BaseHook):
                 "after": tenacity.after_log(log, logging.DEBUG),
             }
             default_kwargs.update(**kwargs)
-            return tenacity.retry(*args, **default_kwargs)(fun)
+            return tenacity.retry(*args, **default_kwargs)(func)
 
         return decorator
 
@@ -442,7 +462,7 @@ class GoogleBaseHook(BaseHook):
     def operation_in_progress_retry(*args, **kwargs) -> Callable[[T], T]:
         """Provide a mechanism to repeat requests in response to operation in progress (HTTP 409) limit."""
 
-        def decorator(fun: T):
+        def decorator(func: T):
             default_kwargs = {
                 "wait": tenacity.wait_exponential(multiplier=1, max=300),
                 "retry": retry_if_operation_in_progress(),
@@ -450,7 +470,25 @@ class GoogleBaseHook(BaseHook):
                 "after": tenacity.after_log(log, logging.DEBUG),
             }
             default_kwargs.update(**kwargs)
-            return cast(T, tenacity.retry(*args, **default_kwargs)(fun))
+            return cast(T, tenacity.retry(*args, **default_kwargs)(func))
+
+        return decorator
+
+    @staticmethod
+    def refresh_credentials_retry(*args, **kwargs) -> Callable[[T], T]:
+        """Provide a mechanism to repeat requests in response to a temporary refresh credential issue."""
+
+        def decorator(func: T):
+            default_kwargs = {
+                "wait": tenacity.wait_exponential(multiplier=1, max=5),
+                "stop": tenacity.stop_after_attempt(3),
+                "retry": retry_if_temporary_refresh_credentials(),
+                "reraise": True,
+                "before": tenacity.before_log(log, logging.DEBUG),
+                "after": tenacity.after_log(log, logging.DEBUG),
+            }
+            default_kwargs.update(**kwargs)
+            return cast(T, tenacity.retry(*args, **default_kwargs)(func))
 
         return decorator
 

airflow/providers/google/firebase/hooks/firestore.py

@@ -25,7 +25,7 @@ from typing import Sequence
 from googleapiclient.discovery import build, build_from_document
 
 from airflow.exceptions import AirflowException
-from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
+from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID, GoogleBaseHook
 
 # Time to sleep between active checks of the operation results
 TIME_TO_SLEEP_IN_SECONDS = 5
@@ -84,7 +84,7 @@ class CloudFirestoreHook(GoogleBaseHook):
 
     @GoogleBaseHook.fallback_to_default_project_id
     def export_documents(
-        self, body: dict, database_id: str = "(default)", project_id: str | None = None
+        self, body: dict, database_id: str = "(default)", project_id: str = PROVIDE_PROJECT_ID
     ) -> None:
         """
         Start a export with the specified configuration.

airflow/providers/google/firebase/operators/firestore.py

@@ -20,6 +20,7 @@ from typing import TYPE_CHECKING, Sequence
 
 from airflow.exceptions import AirflowException
 from airflow.models import BaseOperator
+from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
 from airflow.providers.google.firebase.hooks.firestore import CloudFirestoreHook
 
 if TYPE_CHECKING:
@@ -64,7 +65,7 @@ class CloudFirestoreExportDatabaseOperator(BaseOperator):
         *,
         body: dict,
         database_id: str = "(default)",
-        project_id: str | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
         gcp_conn_id: str = "google_cloud_default",
         api_version: str = "v1",
         impersonation_chain: str | Sequence[str] | None = None,

airflow/providers/google/get_provider_info.py

@@ -28,8 +28,9 @@ def get_provider_info():
         "name": "Google",
         "description": "Google services including:\n\n  - `Google Ads <https://ads.google.com/>`__\n  - `Google Cloud (GCP) <https://cloud.google.com/>`__\n  - `Google Firebase <https://firebase.google.com/>`__\n  - `Google LevelDB <https://github.com/google/leveldb/>`__\n  - `Google Marketing Platform <https://marketingplatform.google.com/>`__\n  - `Google Workspace <https://workspace.google.com/>`__ (formerly Google Suite)\n",
         "state": "ready",
-        "source-date-epoch": 1712665855,
+        "source-date-epoch": 1714476421,
         "versions": [
+            "10.18.0",
             "10.17.0",
             "10.16.0",
             "10.15.0",
@@ -86,7 +87,7 @@ def get_provider_info():
             "1.0.0",
         ],
         "dependencies": [
-            "apache-airflow>=2.6.0",
+            "apache-airflow>=2.7.0",
             "apache-airflow-providers-common-sql>=1.7.2",
             "asgiref>=3.5.2",
             "gcloud-aio-auth>=4.0.0,<5.0.0",

{apache_airflow_providers_google-10.17.0rc1.dist-info → apache_airflow_providers_google-10.18.0rc1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: apache-airflow-providers-google
-Version: 10.17.0rc1
+Version: 10.18.0rc1
 Summary: Provider package apache-airflow-providers-google for Apache Airflow
 Keywords: airflow-provider,google,airflow,integration
 Author-email: Apache Software Foundation <dev@airflow.apache.org>
@@ -23,7 +23,7 @@ Classifier: Programming Language :: Python :: 3.12
 Classifier: Topic :: System :: Monitoring
 Requires-Dist: PyOpenSSL
 Requires-Dist: apache-airflow-providers-common-sql>=1.7.2rc0
-Requires-Dist: apache-airflow>=2.6.0rc0
+Requires-Dist: apache-airflow>=2.7.0rc0
 Requires-Dist: asgiref>=3.5.2
 Requires-Dist: gcloud-aio-auth>=4.0.0,<5.0.0
 Requires-Dist: gcloud-aio-bigquery>=6.1.2
@@ -102,8 +102,8 @@ Requires-Dist: apache-airflow-providers-sftp ; extra == "sftp"
 Requires-Dist: apache-airflow-providers-ssh ; extra == "ssh"
 Requires-Dist: apache-airflow-providers-trino ; extra == "trino"
 Project-URL: Bug Tracker, https://github.com/apache/airflow/issues
-Project-URL: Changelog, https://airflow.apache.org/docs/apache-airflow-providers-google/10.17.0/changelog.html
-Project-URL: Documentation, https://airflow.apache.org/docs/apache-airflow-providers-google/10.17.0
+Project-URL: Changelog, https://airflow.apache.org/docs/apache-airflow-providers-google/10.18.0/changelog.html
+Project-URL: Documentation, https://airflow.apache.org/docs/apache-airflow-providers-google/10.18.0
 Project-URL: Slack Chat, https://s.apache.org/airflow-slack
 Project-URL: Source Code, https://github.com/apache/airflow
 Project-URL: Twitter, https://twitter.com/ApacheAirflow
@@ -171,7 +171,7 @@ Provides-Extra: trino
 
 Package ``apache-airflow-providers-google``
 
-Release: ``10.17.0.rc1``
+Release: ``10.18.0.rc1``
 
 
 Google services including:
@@ -191,7 +191,7 @@ This is a provider package for ``google`` provider. All classes for this provide
 are in ``airflow.providers.google`` python package.
 
 You can find package information and changelog for the provider
-in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-google/10.17.0/>`_.
+in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-google/10.18.0/>`_.
 
 Installation
 ------------
@@ -208,7 +208,7 @@ Requirements
 =======================================  =====================
 PIP package                              Version required
 =======================================  =====================
-``apache-airflow``                       ``>=2.6.0``
+``apache-airflow``                       ``>=2.7.0``
 ``apache-airflow-providers-common-sql``  ``>=1.7.2``
 ``asgiref``                              ``>=3.5.2``
 ``gcloud-aio-auth``                      ``>=4.0.0,<5.0.0``
@@ -307,4 +307,4 @@ Dependent package
 ========================================================================================================================  ====================
 
 The changelog for the provider package can be found in the
-`changelog <https://airflow.apache.org/docs/apache-airflow-providers-google/10.17.0/changelog.html>`_.
+`changelog <https://airflow.apache.org/docs/apache-airflow-providers-google/10.18.0/changelog.html>`_.