apache-airflow-providers-google 10.14.0rc1__py3-none-any.whl → 10.15.0rc1__py3-none-any.whl

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

@@ -19,6 +19,7 @@
 from __future__ import annotations
 
 import asyncio
+import re
 import time
 from typing import Any, AsyncIterator, Sequence
 
@@ -26,6 +27,7 @@ from google.api_core.exceptions import NotFound
 from google.cloud.dataproc_v1 import Batch, ClusterStatus, JobStatus
 
 from airflow.providers.google.cloud.hooks.dataproc import DataprocAsyncHook
+from airflow.providers.google.cloud.utils.dataproc import DataprocOperationType
 from airflow.triggers.base import BaseTrigger, TriggerEvent
 
 
@@ -180,7 +182,7 @@ class DataprocBatchTrigger(DataprocBaseTrigger):
         self.batch_id = batch_id
 
     def serialize(self) -> tuple[str, dict[str, Any]]:
-        """Serializes DataprocBatchTrigger arguments and classpath."""
+        """Serialize DataprocBatchTrigger arguments and classpath."""
         return (
             "airflow.providers.google.cloud.triggers.dataproc.DataprocBatchTrigger",
             {
@@ -242,7 +244,7 @@ class DataprocDeleteClusterTrigger(DataprocBaseTrigger):
         self.metadata = metadata
 
     def serialize(self) -> tuple[str, dict[str, Any]]:
-        """Serializes DataprocDeleteClusterTrigger arguments and classpath."""
+        """Serialize DataprocDeleteClusterTrigger arguments and classpath."""
         return (
             "airflow.providers.google.cloud.triggers.dataproc.DataprocDeleteClusterTrigger",
             {
@@ -281,22 +283,24 @@ class DataprocDeleteClusterTrigger(DataprocBaseTrigger):
             yield TriggerEvent({"status": "error", "message": "Timeout"})
 
 
-class DataprocWorkflowTrigger(DataprocBaseTrigger):
+class DataprocOperationTrigger(DataprocBaseTrigger):
     """
-    Trigger that periodically polls information from Dataproc API to verify status.
+    Trigger that periodically polls information on a long running operation from Dataproc API to verify status.
 
     Implementation leverages asynchronous transport.
     """
 
-    def __init__(self, name: str, **kwargs: Any):
+    def __init__(self, name: str, operation_type: str | None = None, **kwargs: Any):
         super().__init__(**kwargs)
         self.name = name
+        self.operation_type = operation_type
 
     def serialize(self):
         return (
-            "airflow.providers.google.cloud.triggers.dataproc.DataprocWorkflowTrigger",
+            "airflow.providers.google.cloud.triggers.dataproc.DataprocOperationTrigger",
             {
                 "name": self.name,
+                "operation_type": self.operation_type,
                 "project_id": self.project_id,
                 "region": self.region,
                 "gcp_conn_id": self.gcp_conn_id,
@@ -317,14 +321,30 @@ class DataprocWorkflowTrigger(DataprocBaseTrigger):
                     else:
                         status = "success"
                         message = "Operation is successfully ended."
-                    yield TriggerEvent(
-                        {
-                            "operation_name": operation.name,
-                            "operation_done": operation.done,
-                            "status": status,
-                            "message": message,
-                        }
-                    )
+                    if self.operation_type == DataprocOperationType.DIAGNOSE.value:
+                        gcs_regex = rb"gs:\/\/[a-z0-9][a-z0-9_-]{1,61}[a-z0-9_\-\/]*"
+                        gcs_uri_value = operation.response.value
+                        match = re.search(gcs_regex, gcs_uri_value)
+                        if match:
+                            output_uri = match.group(0).decode("utf-8", "ignore")
+                        else:
+                            output_uri = gcs_uri_value
+                        yield TriggerEvent(
+                            {
+                                "status": status,
+                                "message": message,
+                                "output_uri": output_uri,
+                            }
+                        )
+                    else:
+                        yield TriggerEvent(
+                            {
+                                "operation_name": operation.name,
+                                "operation_done": operation.done,
+                                "status": status,
+                                "message": message,
+                            }
+                        )
                     return
                 else:
                     self.log.info("Sleeping for %s seconds.", self.polling_interval_seconds)

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

@@ -60,7 +60,7 @@ class GCSBlobTrigger(BaseTrigger):
         self.hook_params = hook_params
 
     def serialize(self) -> tuple[str, dict[str, Any]]:
-        """Serializes GCSBlobTrigger arguments and classpath."""
+        """Serialize GCSBlobTrigger arguments and classpath."""
         return (
             "airflow.providers.google.cloud.triggers.gcs.GCSBlobTrigger",
             {
@@ -93,7 +93,7 @@ class GCSBlobTrigger(BaseTrigger):
 
     async def _object_exists(self, hook: GCSAsyncHook, bucket_name: str, object_name: str) -> str:
         """
-        Checks for the existence of a file in Google Cloud Storage.
+        Check for the existence of a file in Google Cloud Storage.
 
         :param bucket_name: The Google Cloud Storage bucket where the object is.
         :param object_name: The name of the blob_name to check in the Google cloud
@@ -143,7 +143,7 @@ class GCSCheckBlobUpdateTimeTrigger(BaseTrigger):
         self.hook_params = hook_params
 
     def serialize(self) -> tuple[str, dict[str, Any]]:
-        """Serializes GCSCheckBlobUpdateTimeTrigger arguments and classpath."""
+        """Serialize GCSCheckBlobUpdateTimeTrigger arguments and classpath."""
         return (
             "airflow.providers.google.cloud.triggers.gcs.GCSCheckBlobUpdateTimeTrigger",
             {
@@ -181,7 +181,7 @@ class GCSCheckBlobUpdateTimeTrigger(BaseTrigger):
         self, hook: GCSAsyncHook, bucket_name: str, object_name: str, target_date: datetime
     ) -> tuple[bool, dict[str, Any]]:
         """
-        Checks if the object in the bucket is updated.
+        Check if the object in the bucket is updated.
 
         :param hook: GCSAsyncHook Hook class
         :param bucket_name: The Google Cloud Storage bucket where the object is.
@@ -248,7 +248,7 @@ class GCSPrefixBlobTrigger(GCSBlobTrigger):
         self.prefix = prefix
 
     def serialize(self) -> tuple[str, dict[str, Any]]:
-        """Serializes GCSPrefixBlobTrigger arguments and classpath."""
+        """Serialize GCSPrefixBlobTrigger arguments and classpath."""
         return (
             "airflow.providers.google.cloud.triggers.gcs.GCSPrefixBlobTrigger",
             {
@@ -282,7 +282,7 @@ class GCSPrefixBlobTrigger(GCSBlobTrigger):
 
     async def _list_blobs_with_prefix(self, hook: GCSAsyncHook, bucket_name: str, prefix: str) -> list[str]:
         """
-        Returns names of blobs which match the given prefix for a given bucket.
+        Return names of blobs which match the given prefix for a given bucket.
 
         :param hook: The async hook to use for listing the blobs
         :param bucket_name: The Google Cloud Storage bucket where the object is.
@@ -344,7 +344,7 @@ class GCSUploadSessionTrigger(GCSPrefixBlobTrigger):
         self.last_activity_time: datetime | None = None
 
     def serialize(self) -> tuple[str, dict[str, Any]]:
-        """Serializes GCSUploadSessionTrigger arguments and classpath."""
+        """Serialize GCSUploadSessionTrigger arguments and classpath."""
         return (
             "airflow.providers.google.cloud.triggers.gcs.GCSUploadSessionTrigger",
             {
@@ -377,7 +377,11 @@ class GCSUploadSessionTrigger(GCSPrefixBlobTrigger):
             yield TriggerEvent({"status": "error", "message": str(e)})
 
     def _get_time(self) -> datetime:
-        """This is just a wrapper of datetime.datetime.now to simplify mocking in the unittests."""
+        """
+        Get current local date and time.
+
+        This is just a wrapper of datetime.datetime.now to simplify mocking in the unittests.
+        """
         return datetime.now()
 
     def _is_bucket_updated(self, current_objects: set[str]) -> dict[str, str]:

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

@@ -174,7 +174,7 @@ class GKEOperationTrigger(BaseTrigger):
         self._hook: GKEAsyncHook | None = None
 
     def serialize(self) -> tuple[str, dict[str, Any]]:
-        """Serializes GKEOperationTrigger arguments and classpath."""
+        """Serialize GKEOperationTrigger arguments and classpath."""
         return (
             "airflow.providers.google.cloud.triggers.kubernetes_engine.GKEOperationTrigger",
             {
@@ -188,7 +188,7 @@ class GKEOperationTrigger(BaseTrigger):
         )
 
     async def run(self) -> AsyncIterator[TriggerEvent]:  # type: ignore[override]
-        """Gets operation status and yields corresponding event."""
+        """Get operation status and yields corresponding event."""
         hook = self._get_hook()
         try:
             while True:

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

@@ -69,7 +69,7 @@ class MLEngineStartTrainingJobTrigger(BaseTrigger):
         self.impersonation_chain = impersonation_chain
 
     def serialize(self) -> tuple[str, dict[str, Any]]:
-        """Serializes MLEngineStartTrainingJobTrigger arguments and classpath."""
+        """Serialize MLEngineStartTrainingJobTrigger arguments and classpath."""
         return (
             "airflow.providers.google.cloud.triggers.mlengine.MLEngineStartTrainingJobTrigger",
             {
@@ -89,7 +89,7 @@ class MLEngineStartTrainingJobTrigger(BaseTrigger):
         )
 
     async def run(self) -> AsyncIterator[TriggerEvent]:  # type: ignore[override]
-        """Gets current job execution status and yields a TriggerEvent."""
+        """Get current job execution status and yields a TriggerEvent."""
         hook = self._get_async_hook()
         try:
             while True:

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

@@ -79,7 +79,7 @@ class PubsubPullTrigger(BaseTrigger):
         self.hook = PubSubAsyncHook()
 
     def serialize(self) -> tuple[str, dict[str, Any]]:
-        """Serializes PubsubPullTrigger arguments and classpath."""
+        """Serialize PubsubPullTrigger arguments and classpath."""
         return (
             "airflow.providers.google.cloud.triggers.pubsub.PubsubPullTrigger",
             {

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

@@ -0,0 +1,99 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import Any, AsyncIterator, Sequence
+
+from google.cloud.aiplatform_v1 import HyperparameterTuningJob, JobState
+
+from airflow.exceptions import AirflowException
+from airflow.providers.google.cloud.hooks.vertex_ai.hyperparameter_tuning_job import (
+    HyperparameterTuningJobAsyncHook,
+)
+from airflow.triggers.base import BaseTrigger, TriggerEvent
+
+
+class CreateHyperparameterTuningJobTrigger(BaseTrigger):
+    """CreateHyperparameterTuningJobTrigger run on the trigger worker to perform create operation."""
+
+    statuses_success = {
+        JobState.JOB_STATE_PAUSED,
+        JobState.JOB_STATE_SUCCEEDED,
+    }
+
+    def __init__(
+        self,
+        conn_id: str,
+        project_id: str,
+        location: str,
+        job_id: str,
+        poll_interval: int,
+        impersonation_chain: str | Sequence[str] | None = None,
+    ):
+        super().__init__()
+        self.conn_id = conn_id
+        self.project_id = project_id
+        self.location = location
+        self.job_id = job_id
+        self.poll_interval = poll_interval
+        self.impersonation_chain = impersonation_chain
+
+    def serialize(self) -> tuple[str, dict[str, Any]]:
+        return (
+            "airflow.providers.google.cloud.triggers.vertex_ai.CreateHyperparameterTuningJobTrigger",
+            {
+                "conn_id": self.conn_id,
+                "project_id": self.project_id,
+                "location": self.location,
+                "job_id": self.job_id,
+                "poll_interval": self.poll_interval,
+                "impersonation_chain": self.impersonation_chain,
+            },
+        )
+
+    async def run(self) -> AsyncIterator[TriggerEvent]:
+        hook = self._get_async_hook()
+        try:
+            job = await hook.wait_hyperparameter_tuning_job(
+                project_id=self.project_id,
+                location=self.location,
+                job_id=self.job_id,
+                poll_interval=self.poll_interval,
+            )
+        except AirflowException as ex:
+            yield TriggerEvent(
+                {
+                    "status": "error",
+                    "message": str(ex),
+                }
+            )
+            return
+
+        status = "success" if job.state in self.statuses_success else "error"
+        message = f"Hyperparameter tuning job {job.name} completed with status {job.state.name}"
+        yield TriggerEvent(
+            {
+                "status": status,
+                "message": message,
+                "job": HyperparameterTuningJob.to_dict(job),
+            }
+        )
+
+    def _get_async_hook(self) -> HyperparameterTuningJobAsyncHook:
+        return HyperparameterTuningJobAsyncHook(
+            gcp_conn_id=self.conn_id, impersonation_chain=self.impersonation_chain
+        )

airflow/providers/google/cloud/utils/bigquery.py

@@ -21,7 +21,7 @@ from typing import Any
 
 def bq_cast(string_field: str, bq_type: str) -> None | int | float | bool | str:
     """
-    Helper method that casts a BigQuery row to the appropriate data types.
+    Cast a BigQuery row to the appropriate data types.
 
     This is useful because BigQuery returns all fields as strings.
     """
@@ -41,7 +41,7 @@ def bq_cast(string_field: str, bq_type: str) -> None | int | float | bool | str:
 
 def convert_job_id(job_id: str | list[str], project_id: str, location: str | None) -> Any:
     """
-    Helper method that converts to path: project_id:location:job_id.
+    Convert job_id to path: project_id:location:job_id.
 
     :param project_id: Required. The ID of the Google Cloud project where workspace located.
     :param location: Optional. The ID of the Google Cloud region where workspace located.

airflow/providers/google/cloud/utils/credentials_provider.py

@@ -358,7 +358,7 @@ class _CredentialProvider(LoggingMixin):
 
 
 def get_credentials_and_project_id(*args, **kwargs) -> tuple[google.auth.credentials.Credentials, str]:
-    """Returns the Credentials object for Google API and the associated project_id."""
+    """Return the Credentials object for Google API and the associated project_id."""
     return _CredentialProvider(*args, **kwargs).get_credentials_and_project()
 
 
@@ -398,7 +398,7 @@ def _get_target_principal_and_delegates(
 
 def _get_project_id_from_service_account_email(service_account_email: str) -> str:
     """
-    Extracts project_id from service account's email address.
+    Extract project_id from service account's email address.
 
     :param service_account_email: email of the service account.
 

airflow/providers/google/cloud/utils/dataform.py

@@ -45,7 +45,7 @@ def make_initialization_workspace_flow(
     without_installation: bool = False,
 ) -> tuple:
     """
-    Creates flow which simulates the initialization of the default project.
+    Create flow which simulates the initialization of the default project.
 
     :param project_id: Required. The ID of the Google Cloud project where workspace located.
     :param region: Required. The ID of the Google Cloud region where workspace located.

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

@@ -0,0 +1,25 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from enum import Enum
+
+
+class DataprocOperationType(Enum):
+    """Contains types of long running operations."""
+
+    DIAGNOSE = "DIAGNOSE"

airflow/providers/google/cloud/utils/field_validator.py

@@ -309,7 +309,7 @@ class GcpBodyFieldValidator(LoggingMixin):
 
     def _validate_field(self, validation_spec, dictionary_to_validate, parent=None, force_optional=False):
         """
-        Validates if field is OK.
+        Validate if field is OK.
 
         :param validation_spec: specification of the field
         :param dictionary_to_validate: dictionary where the field should be present
@@ -413,7 +413,7 @@ class GcpBodyFieldValidator(LoggingMixin):
 
     def validate(self, body_to_validate: dict) -> None:
         """
-        Validates if the body (dictionary) follows specification that the validator was instantiated with.
+        Validate if the body (dictionary) follows specification that the validator was instantiated with.
 
         Raises ValidationSpecificationException or ValidationFieldException in case of problems
         with specification or the body not conforming to the specification respectively.

airflow/providers/google/cloud/utils/helpers.py

@@ -19,12 +19,12 @@ from __future__ import annotations
 
 
 def normalize_directory_path(source_object: str | None) -> str | None:
-    """Makes sure dir path ends with a slash."""
+    """Make sure dir path ends with a slash."""
     return source_object + "/" if source_object and not source_object.endswith("/") else source_object
 
 
 def resource_path_to_dict(resource_name: str) -> dict[str, str]:
-    """Converts a path-like GCP resource name into a dictionary.
+    """Convert a path-like GCP resource name into a dictionary.
 
     For example, the path `projects/my-project/locations/my-location/instances/my-instance` will be converted
     to a dict:

airflow/providers/google/cloud/utils/mlengine_operator_utils.py

@@ -57,7 +57,7 @@ def create_evaluate_ops(
     py_interpreter="python3",
 ) -> tuple[MLEngineStartBatchPredictionJobOperator, BeamRunPythonPipelineOperator, PythonOperator]:
     r"""
-    Creates Operators needed for model evaluation and returns.
+    Create Operators needed for model evaluation and returns.
 
     This function is deprecated. All the functionality of legacy MLEngine and new features are available
     on the Vertex AI platform.

airflow/providers/google/cloud/utils/mlengine_prediction_summary.py

@@ -151,7 +151,7 @@ def MakeSummary(pcoll, metric_fn, metric_keys):
 
 
 def run(argv=None):
-    """Helper for obtaining prediction summary."""
+    """Obtain prediction summary."""
     parser = argparse.ArgumentParser()
     parser.add_argument(
         "--prediction_path",

airflow/providers/google/common/auth_backend/google_openid.py

@@ -52,7 +52,7 @@ def create_client_session():
 
 
 def init_app(_):
-    """Initializes authentication."""
+    """Initialize authentication."""
 
 
 def _get_id_token_from_request(request) -> str | None:
@@ -110,7 +110,7 @@ T = TypeVar("T", bound=Callable)
 
 
 def requires_authentication(function: T):
-    """Decorator for functions that require authentication."""
+    """Decorator for function that require authentication."""
 
     @wraps(function)
     def decorated(*args, **kwargs):