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

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

@@ -16,19 +16,44 @@
 # under the License.
 from __future__ import annotations
 
-from typing import Any, AsyncIterator, Sequence
+from functools import cached_property
+from typing import TYPE_CHECKING, Any, AsyncIterator, Sequence
 
-from google.cloud.aiplatform_v1 import HyperparameterTuningJob, JobState
+from google.cloud.aiplatform_v1 import (
+    BatchPredictionJob,
+    HyperparameterTuningJob,
+    JobState,
+    PipelineState,
+    types,
+)
 
 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.hyperparameter_tuning_job import (
     HyperparameterTuningJobAsyncHook,
 )
+from airflow.providers.google.cloud.hooks.vertex_ai.pipeline_job import PipelineJobAsyncHook
 from airflow.triggers.base import BaseTrigger, TriggerEvent
 
+if TYPE_CHECKING:
+    from proto import Message
 
-class CreateHyperparameterTuningJobTrigger(BaseTrigger):
-    """CreateHyperparameterTuningJobTrigger run on the trigger worker to perform create operation."""
+
+class BaseVertexAIJobTrigger(BaseTrigger):
+    """Base class for Vertex AI job triggers.
+
+    This trigger polls the Vertex AI job and checks its status.
+
+    In order to use it properly, you must:
+    - implement the following methods `_wait_job()`.
+    - override required `job_type_verbose_name` attribute to provide meaningful message describing your
+    job type.
+    - override required `job_serializer_class` attribute to provide proto.Message class that will be used
+    to serialize your job with `to_dict()` class method.
+    """
+
+    job_type_verbose_name: str = "Vertex AI Job"
+    job_serializer_class: Message = None
 
     statuses_success = {
         JobState.JOB_STATE_PAUSED,
@@ -51,10 +76,13 @@ class CreateHyperparameterTuningJobTrigger(BaseTrigger):
         self.job_id = job_id
         self.poll_interval = poll_interval
         self.impersonation_chain = impersonation_chain
+        self.trigger_class_path = (
+            f"airflow.providers.google.cloud.triggers.vertex_ai.{self.__class__.__name__}"
+        )
 
     def serialize(self) -> tuple[str, dict[str, Any]]:
         return (
-            "airflow.providers.google.cloud.triggers.vertex_ai.CreateHyperparameterTuningJobTrigger",
+            self.trigger_class_path,
             {
                 "conn_id": self.conn_id,
                 "project_id": self.project_id,
@@ -66,14 +94,8 @@ class CreateHyperparameterTuningJobTrigger(BaseTrigger):
         )
 
     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,
-            )
+            job = await self._wait_job()
         except AirflowException as ex:
             yield TriggerEvent(
                 {
@@ -84,16 +106,86 @@ class CreateHyperparameterTuningJobTrigger(BaseTrigger):
             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}"
+        message = f"{self.job_type_verbose_name} {job.name} completed with status {job.state.name}"
         yield TriggerEvent(
             {
                 "status": status,
                 "message": message,
-                "job": HyperparameterTuningJob.to_dict(job),
+                "job": self._serialize_job(job),
             }
         )
 
-    def _get_async_hook(self) -> HyperparameterTuningJobAsyncHook:
+    async def _wait_job(self) -> Any:
+        """Awaits a Vertex AI job instance for a status examination."""
+        raise NotImplementedError
+
+    def _serialize_job(self, job: Any) -> Any:
+        return self.job_serializer_class.to_dict(job)
+
+
+class CreateHyperparameterTuningJobTrigger(BaseVertexAIJobTrigger):
+    """CreateHyperparameterTuningJobTrigger run on the trigger worker to perform create operation."""
+
+    job_type_verbose_name = "Hyperparameter Tuning Job"
+    job_serializer_class = HyperparameterTuningJob
+
+    @cached_property
+    def async_hook(self) -> HyperparameterTuningJobAsyncHook:
         return HyperparameterTuningJobAsyncHook(
             gcp_conn_id=self.conn_id, impersonation_chain=self.impersonation_chain
         )
+
+    async def _wait_job(self) -> types.HyperparameterTuningJob:
+        job: types.HyperparameterTuningJob = await self.async_hook.wait_hyperparameter_tuning_job(
+            project_id=self.project_id,
+            location=self.location,
+            job_id=self.job_id,
+            poll_interval=self.poll_interval,
+        )
+        return job
+
+
+class CreateBatchPredictionJobTrigger(BaseVertexAIJobTrigger):
+    """CreateBatchPredictionJobTrigger run on the trigger worker to perform create operation."""
+
+    job_type_verbose_name = "Batch Prediction Job"
+    job_serializer_class = BatchPredictionJob
+
+    @cached_property
+    def async_hook(self) -> BatchPredictionJobAsyncHook:
+        return BatchPredictionJobAsyncHook(
+            gcp_conn_id=self.conn_id, impersonation_chain=self.impersonation_chain
+        )
+
+    async def _wait_job(self) -> types.BatchPredictionJob:
+        job: types.BatchPredictionJob = await self.async_hook.wait_batch_prediction_job(
+            project_id=self.project_id,
+            location=self.location,
+            job_id=self.job_id,
+            poll_interval=self.poll_interval,
+        )
+        return job
+
+
+class RunPipelineJobTrigger(BaseVertexAIJobTrigger):
+    """Make async calls to Vertex AI to check the state of a Pipeline Job."""
+
+    job_type_verbose_name = "Pipeline Job"
+    job_serializer_class = types.PipelineJob
+    statuses_success = {
+        PipelineState.PIPELINE_STATE_PAUSED,
+        PipelineState.PIPELINE_STATE_SUCCEEDED,
+    }
+
+    @cached_property
+    def async_hook(self) -> PipelineJobAsyncHook:
+        return PipelineJobAsyncHook(gcp_conn_id=self.conn_id, impersonation_chain=self.impersonation_chain)
+
+    async def _wait_job(self) -> types.PipelineJob:
+        job: types.PipelineJob = await self.async_hook.wait_for_pipeline_job(
+            project_id=self.project_id,
+            location=self.location,
+            job_id=self.job_id,
+            poll_interval=self.poll_interval,
+        )
+        return job

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

@@ -67,7 +67,7 @@ specification of the path you should delete - separated with '.'
 >>>         },
 >>>     }
 >>> }
->>> sanitizer=GcpBodyFieldSanitizer(FIELDS_TO_SANITIZE)
+>>> sanitizer = GcpBodyFieldSanitizer(FIELDS_TO_SANITIZE)
 >>> sanitizer.sanitize(body)
 >>> json.dumps(body, indent=2)
 {
@@ -96,6 +96,7 @@ In case  they are dictionaries, subsequent component names key of the field, in
 arrays - the sanitizer iterates through all dictionaries in the array and searches
 components in all elements of the array.
 """
+
 from __future__ import annotations
 
 from airflow.exceptions import AirflowException

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

@@ -129,6 +129,7 @@ Here are the guidelines that you should follow to make validation forward-compat
   backwards-incompatible changes that might sometimes occur in the APIs.
 
 """
+
 from __future__ import annotations
 
 import re

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

@@ -15,6 +15,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """This module contains helper functions for Google Cloud operators."""
+
 from __future__ import annotations
 
 

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

@@ -15,6 +15,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """This module contains helper functions for MLEngine operators."""
+
 from __future__ import annotations
 
 import base64

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

@@ -106,6 +106,7 @@ To test outside of the dag:
 
     pcoll
 """
+
 from __future__ import annotations
 
 import argparse

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

@@ -16,6 +16,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """This module contains code related to OpenLineage and lineage extraction."""
+
 from __future__ import annotations
 
 from typing import TYPE_CHECKING, Any

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

@@ -16,6 +16,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """Authentication backend that use Google credentials for authorization."""
+
 from __future__ import annotations
 
 import logging

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

@@ -16,6 +16,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """This module contains a Google Cloud API base hook."""
+
 from __future__ import annotations
 
 import datetime
@@ -316,7 +317,7 @@ class GoogleBaseHook(BaseHook):
         credentials.refresh(auth_req)
         return credentials.token
 
-    @functools.lru_cache(maxsize=None)
+    @functools.cached_property
     def _get_credentials_email(self) -> str:
         """
         Return the email address associated with the currently logged in account.

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

@@ -16,6 +16,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """This module allows you to connect to the Google Discovery API Service and query it."""
+
 from __future__ import annotations
 
 from typing import Sequence

airflow/providers/google/common/links/storage.py

@@ -15,6 +15,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """This module contains a link for GCS Storage assets."""
+
 from __future__ import annotations
 
 from typing import TYPE_CHECKING

airflow/providers/google/common/utils/id_token_credentials.py

@@ -28,6 +28,7 @@ To obtain info about this token, run the following commands:
 
     RefreshError
 """
+
 from __future__ import annotations
 
 import json

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

@@ -16,6 +16,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """Hook for Google Cloud Firestore service."""
+
 from __future__ import annotations
 
 import time

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": 1709555588,
+        "source-date-epoch": 1712665855,
         "versions": [
+            "10.17.0",
             "10.16.0",
             "10.15.0",
             "10.14.0",
@@ -92,14 +93,15 @@ def get_provider_info():
             "gcloud-aio-bigquery>=6.1.2",
             "gcloud-aio-storage>=9.0.0",
             "gcsfs>=2023.10.0",
-            "google-ads>=22.1.0",
+            "google-ads>=23.1.0",
             "google-analytics-admin",
             "google-api-core>=2.11.0,!=2.16.0",
             "google-api-python-client>=1.6.0",
             "google-auth>=1.0.0",
             "google-auth-httplib2>=0.0.1",
-            "google-cloud-aiplatform>=1.22.1",
+            "google-cloud-aiplatform>=1.42.1",
             "google-cloud-automl>=2.12.0",
+            "google-cloud-bigquery>=3.0.1",
             "google-cloud-bigquery-datatransfer>=3.13.0",
             "google-cloud-bigtable>=2.17.0",
             "google-cloud-build>=3.22.0",
@@ -144,6 +146,7 @@ def get_provider_info():
             "PyOpenSSL",
             "sqlalchemy-bigquery>=1.2.1",
             "sqlalchemy-spanner>=1.6.2",
+            "python-slugify>=5.0",
         ],
         "additional-extras": [
             {"name": "apache.beam", "dependencies": ["apache-beam[gcp]"]},
@@ -807,6 +810,7 @@ def get_provider_info():
                     "airflow.providers.google.cloud.operators.vertex_ai.hyperparameter_tuning_job",
                     "airflow.providers.google.cloud.operators.vertex_ai.model_service",
                     "airflow.providers.google.cloud.operators.vertex_ai.pipeline_job",
+                    "airflow.providers.google.cloud.operators.vertex_ai.generative_model",
                 ],
             },
             {
@@ -1141,6 +1145,7 @@ def get_provider_info():
                     "airflow.providers.google.cloud.hooks.vertex_ai.hyperparameter_tuning_job",
                     "airflow.providers.google.cloud.hooks.vertex_ai.model_service",
                     "airflow.providers.google.cloud.hooks.vertex_ai.pipeline_job",
+                    "airflow.providers.google.cloud.hooks.vertex_ai.generative_model",
                 ],
             },
             {
@@ -1532,6 +1537,7 @@ def get_provider_info():
             "airflow.providers.google.cloud.links.kubernetes_engine.KubernetesEngineClusterLink",
             "airflow.providers.google.cloud.links.kubernetes_engine.KubernetesEnginePodLink",
             "airflow.providers.google.cloud.links.kubernetes_engine.KubernetesEngineJobLink",
+            "airflow.providers.google.cloud.links.kubernetes_engine.KubernetesEngineWorkloadsLink",
             "airflow.providers.google.cloud.links.pubsub.PubSubSubscriptionLink",
             "airflow.providers.google.cloud.links.pubsub.PubSubTopicLink",
             "airflow.providers.google.cloud.links.cloud_memorystore.MemcachedInstanceDetailsLink",

airflow/providers/google/go_module_utils.py

@@ -16,6 +16,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """Utilities initializing and managing Go modules."""
+
 from __future__ import annotations
 
 import os

airflow/providers/google/leveldb/hooks/leveldb.py

@@ -15,6 +15,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """Hook for Level DB."""
+
 from __future__ import annotations
 
 from airflow.exceptions import AirflowException, AirflowOptionalProviderFeatureException
@@ -93,7 +94,7 @@ class LevelDBHook(BaseHook):
         """
         if command == "put":
             if not value:
-                raise Exception("Please provide `value`!")
+                raise ValueError("Please provide `value`!")
             return self.put(key, value)
         elif command == "get":
             return self.get(key)
@@ -101,9 +102,9 @@ class LevelDBHook(BaseHook):
             return self.delete(key)
         elif command == "write_batch":
             if not keys:
-                raise Exception("Please provide `keys`!")
+                raise ValueError("Please provide `keys`!")
             if not values:
-                raise Exception("Please provide `values`!")
+                raise ValueError("Please provide `values`!")
             return self.write_batch(keys, values)
         else:
             raise LevelDBHookException("Unknown command for LevelDB hook")
@@ -116,7 +117,7 @@ class LevelDBHook(BaseHook):
         :param value: value for put execution e.g. ``b'value'``, ``b'another-value'``
         """
         if not self.db:
-            raise Exception(DB_NOT_INITIALIZED_BEFORE)
+            raise AirflowException(DB_NOT_INITIALIZED_BEFORE)
         self.db.put(key, value)
 
     def get(self, key: bytes) -> bytes:
@@ -127,7 +128,7 @@ class LevelDBHook(BaseHook):
         :returns: value of key from db.get
         """
         if not self.db:
-            raise Exception(DB_NOT_INITIALIZED_BEFORE)
+            raise AirflowException(DB_NOT_INITIALIZED_BEFORE)
         return self.db.get(key)
 
     def delete(self, key: bytes):
@@ -137,7 +138,7 @@ class LevelDBHook(BaseHook):
         :param key: key for delete execution, e.g. ``b'key'``, ``b'another-key'``
         """
         if not self.db:
-            raise Exception(DB_NOT_INITIALIZED_BEFORE)
+            raise AirflowException(DB_NOT_INITIALIZED_BEFORE)
         self.db.delete(key)
 
     def write_batch(self, keys: list[bytes], values: list[bytes]):
@@ -148,7 +149,7 @@ class LevelDBHook(BaseHook):
         :param values: values for write_batch execution e.g. ``[b'value', b'another-value']``
         """
         if not self.db:
-            raise Exception(DB_NOT_INITIALIZED_BEFORE)
+            raise AirflowException(DB_NOT_INITIALIZED_BEFORE)
         with self.db.write_batch() as batch:
             for i, key in enumerate(keys):
                 batch.put(key, values[i])

airflow/providers/google/marketing_platform/example_dags/example_display_video.py

@@ -18,6 +18,7 @@
 """
 Example Airflow DAG that shows how to use DisplayVideo.
 """
+
 from __future__ import annotations
 
 import os

airflow/providers/google/marketing_platform/hooks/analytics_admin.py

@@ -24,6 +24,7 @@ Hooks for Google Analytics (GA4) Admin service.
     ListAccountsPager
     ListGoogleAdsLinksPager
 """
+
 from __future__ import annotations
 
 from typing import TYPE_CHECKING, Sequence

airflow/providers/google/marketing_platform/hooks/campaign_manager.py

@@ -16,6 +16,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """This module contains Google Campaign Manager hook."""
+
 from __future__ import annotations
 
 from typing import TYPE_CHECKING, Any, Sequence

airflow/providers/google/marketing_platform/hooks/display_video.py

@@ -16,6 +16,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """This module contains Google DisplayVideo hook."""
+
 from __future__ import annotations
 
 from typing import Any, Sequence

airflow/providers/google/marketing_platform/hooks/search_ads.py

@@ -16,6 +16,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """This module contains Google Search Ads 360 hook."""
+
 from __future__ import annotations
 
 from typing import Any, Sequence

airflow/providers/google/marketing_platform/operators/analytics.py

@@ -16,6 +16,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """This module contains Google Analytics 360 operators."""
+
 from __future__ import annotations
 
 import csv

airflow/providers/google/marketing_platform/operators/analytics_admin.py

@@ -16,6 +16,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """This module contains Google Analytics 4 (GA4) operators."""
+
 from __future__ import annotations
 
 from typing import TYPE_CHECKING, Any, Sequence
@@ -105,8 +106,9 @@ class GoogleAnalyticsAdminListAccountsOperator(GoogleCloudBaseOperator):
             impersonation_chain=self.impersonation_chain,
         )
         self.log.info(
-            "Requesting list of Google Analytics accounts. "
-            f"Page size: {self.page_size}, page token: {self.page_token}"
+            "Requesting list of Google Analytics accounts. Page size: %s, page token: %s",
+            self.page_size,
+            self.page_token,
         )
         accounts = hook.list_accounts(
             page_size=self.page_size,

airflow/providers/google/marketing_platform/operators/campaign_manager.py

@@ -16,6 +16,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """This module contains Google CampaignManager operators."""
+
 from __future__ import annotations
 
 import json

airflow/providers/google/marketing_platform/operators/display_video.py

@@ -16,6 +16,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """This module contains Google DisplayVideo operators."""
+
 from __future__ import annotations
 
 import csv

airflow/providers/google/marketing_platform/operators/search_ads.py

@@ -16,6 +16,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """This module contains Google Search Ads operators."""
+
 from __future__ import annotations
 
 import json

airflow/providers/google/marketing_platform/sensors/campaign_manager.py

@@ -16,6 +16,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """This module contains Google Campaign Manager sensor."""
+
 from __future__ import annotations
 
 from typing import TYPE_CHECKING, Sequence

airflow/providers/google/marketing_platform/sensors/display_video.py

@@ -15,6 +15,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """Sensor for detecting the completion of DV360 reports."""
+
 from __future__ import annotations
 
 from typing import TYPE_CHECKING, Sequence
@@ -157,7 +158,7 @@ class GoogleDisplayVideo360RunQuerySensor(BaseSensorOperator):
 
         response = hook.get_report(query_id=self.query_id, report_id=self.report_id)
         status = response.get("metadata", {}).get("status", {}).get("state")
-        self.log.info(f"STATUS OF THE REPORT {self.report_id} FOR QUERY {self.query_id}: {status}")
+        self.log.info("STATUS OF THE REPORT %s FOR QUERY %s: %s", self.report_id, self.query_id, status)
         if response and status in ["DONE", "FAILED"]:
             return True
         return False

airflow/providers/google/marketing_platform/sensors/search_ads.py

@@ -16,6 +16,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """This module contains Google Search Ads sensor."""
+
 from __future__ import annotations
 
 from typing import TYPE_CHECKING, Sequence

airflow/providers/google/suite/hooks/calendar.py

@@ -16,6 +16,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """This module contains a Google Calendar API hook."""
+
 from __future__ import annotations
 
 from typing import TYPE_CHECKING, Any, Sequence

airflow/providers/google/suite/hooks/drive.py

@@ -16,6 +16,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """Hook for Google Drive service."""
+
 from __future__ import annotations
 
 from typing import IO, Any, Sequence

airflow/providers/google/suite/hooks/sheets.py

@@ -16,6 +16,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """This module contains a Google Sheets API hook."""
+
 from __future__ import annotations
 
 from typing import Any, Sequence

airflow/providers/google/suite/sensors/drive.py

@@ -16,6 +16,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """This module contains Google Drive sensors."""
+
 from __future__ import annotations
 
 from typing import TYPE_CHECKING, Sequence

airflow/providers/google/suite/transfers/gcs_to_gdrive.py

@@ -16,6 +16,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """This module contains a Google Cloud Storage to Google Drive transfer operator."""
+
 from __future__ import annotations
 
 import tempfile
@@ -79,6 +80,9 @@ class GCSToGoogleDriveOperator(BaseOperator):
         If set as a sequence, the identities from the list must grant
         Service Account Token Creator IAM role to the directly preceding identity, with first
         account from the list granting this role to the originating account (templated).
+    :param delegate_to: (Optional) The account to impersonate using domain-wide delegation
+        of authority, if any. For this to work, the service account making the
+        request must have domain-wide delegation enabled. This only applies to the Google Drive connection.
     """
 
     template_fields: Sequence[str] = (
@@ -99,6 +103,7 @@ class GCSToGoogleDriveOperator(BaseOperator):
         move_object: bool = False,
         gcp_conn_id: str = "google_cloud_default",
         impersonation_chain: str | Sequence[str] | None = None,
+        delegate_to: str | None = None,
         **kwargs,
     ) -> None:
         super().__init__(**kwargs)
@@ -110,6 +115,7 @@ class GCSToGoogleDriveOperator(BaseOperator):
         self.move_object = move_object
         self.gcp_conn_id = gcp_conn_id
         self.impersonation_chain = impersonation_chain
+        self.delegate_to = delegate_to
         self.gcs_hook: GCSHook | None = None
         self.gdrive_hook: GoogleDriveHook | None = None
 
@@ -121,6 +127,7 @@ class GCSToGoogleDriveOperator(BaseOperator):
         self.gdrive_hook = GoogleDriveHook(
             gcp_conn_id=self.gcp_conn_id,
             impersonation_chain=self.impersonation_chain,
+            delegate_to=self.delegate_to,
         )
 
         if WILDCARD in self.source_object: