apache-airflow-providers-google 10.19.0__py3-none-any.whl → 10.20.0__py3-none-any.whl

airflow/providers/google/LICENSE

@@ -215,7 +215,7 @@ Third party Apache 2.0 licenses
 
 The following components are provided under the Apache 2.0 License.
 See project link for details. The text of each license is also included
-at licenses/LICENSE-[project].txt.
+at 3rd-party-licenses/LICENSE-[project].txt.
 
     (ALv2 License) hue v4.3.0 (https://github.com/cloudera/hue/)
     (ALv2 License) jqclock v2.3.0 (https://github.com/JohnRDOrazio/jQuery-Clock-Plugin)
@@ -227,7 +227,7 @@ MIT licenses
 ========================================================================
 
 The following components are provided under the MIT License. See project link for details.
-The text of each license is also included at licenses/LICENSE-[project].txt.
+The text of each license is also included at 3rd-party-licenses/LICENSE-[project].txt.
 
     (MIT License) jquery v3.5.1 (https://jquery.org/license/)
     (MIT License) dagre-d3 v0.6.4 (https://github.com/cpettitt/dagre-d3)
@@ -243,11 +243,11 @@ The text of each license is also included at licenses/LICENSE-[project].txt.
 BSD 3-Clause licenses
 ========================================================================
 The following components are provided under the BSD 3-Clause license. See project links for details.
-The text of each license is also included at licenses/LICENSE-[project].txt.
+The text of each license is also included at 3rd-party-licenses/LICENSE-[project].txt.
 
     (BSD 3 License) d3 v5.16.0 (https://d3js.org)
     (BSD 3 License) d3-shape v2.1.0 (https://github.com/d3/d3-shape)
     (BSD 3 License) cgroupspy 0.2.1 (https://github.com/cloudsigma/cgroupspy)
 
 ========================================================================
-See licenses/LICENSES-ui.txt for packages used in `/airflow/www`
+See 3rd-party-licenses/LICENSES-ui.txt for packages used in `/airflow/www`

airflow/providers/google/__init__.py

@@ -29,7 +29,7 @@ from airflow import __version__ as airflow_version
 
 __all__ = ["__version__"]
 
-__version__ = "10.19.0"
+__version__ = "10.20.0"
 
 if packaging.version.parse(packaging.version.parse(airflow_version).base_version) < packaging.version.parse(
     "2.7.0"

airflow/providers/google/ads/hooks/ads.py

@@ -32,9 +32,9 @@ from airflow.hooks.base import BaseHook
 from airflow.providers.google.common.hooks.base_google import get_field
 
 if TYPE_CHECKING:
-    from google.ads.googleads.v16.services.services.customer_service import CustomerServiceClient
-    from google.ads.googleads.v16.services.services.google_ads_service import GoogleAdsServiceClient
-    from google.ads.googleads.v16.services.types.google_ads_service import GoogleAdsRow
+    from google.ads.googleads.v17.services.services.customer_service import CustomerServiceClient
+    from google.ads.googleads.v17.services.services.google_ads_service import GoogleAdsServiceClient
+    from google.ads.googleads.v17.services.types.google_ads_service import GoogleAdsRow
     from google.api_core.page_iterator import GRPCIterator
 
 
@@ -100,7 +100,7 @@ class GoogleAdsHook(BaseHook):
     :param api_version: The Google Ads API version to use.
     """
 
-    default_api_version = "v16"
+    default_api_version = "v17"
 
     def __init__(
         self,

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

@@ -344,6 +344,32 @@ class CloudDataTransferServiceHook(GoogleBaseHook):
             .execute(num_retries=self.num_retries)
         )
 
+    @GoogleBaseHook.fallback_to_default_project_id
+    def run_transfer_job(self, job_name: str, project_id: str) -> dict:
+        """Run Google Storage Transfer Service job.
+
+        :param job_name: (Required) Name of the job to be fetched
+        :param project_id: (Optional) the ID of the project that owns the Transfer
+            Job. If set to None or missing, the default project_id from the Google Cloud
+            connection is used.
+        :return: If successful, Operation. See:
+            https://cloud.google.com/storage-transfer/docs/reference/rest/v1/Operation
+
+        .. seealso:: https://cloud.google.com/storage-transfer/docs/reference/rest/v1/transferJobs/run
+
+        """
+        return (
+            self.get_conn()
+            .transferJobs()
+            .run(
+                jobName=job_name,
+                body={
+                    PROJECT_ID: project_id,
+                },
+            )
+            .execute(num_retries=self.num_retries)
+        )
+
     def cancel_transfer_operation(self, operation_name: str) -> None:
         """Cancel a transfer operation in Google Storage Transfer Service.
 

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

@@ -71,7 +71,7 @@ DEFAULT_DATAFLOW_LOCATION = "us-central1"
 
 
 JOB_ID_PATTERN = re.compile(
-    r"Submitted job: (?P<job_id_java>.*)|Created job with id: \[(?P<job_id_python>.*)\]"
+    r"Submitted job: (?P<job_id_java>[^\"\n]*)|Created job with id: \[(?P<job_id_python>[^\"\n]*)\]"
 )
 
 T = TypeVar("T", bound=Callable)
@@ -582,6 +582,11 @@ class DataflowHook(GoogleBaseHook):
         http_authorized = self._authorize()
         return build("dataflow", "v1b3", http=http_authorized, cache_discovery=False)
 
+    def get_pipelines_conn(self) -> build:
+        """Return a Google Cloud Data Pipelines service object."""
+        http_authorized = self._authorize()
+        return build("datapipelines", "v1", http=http_authorized, cache_discovery=False)
+
     @_fallback_to_location_from_variables
     @_fallback_to_project_id_from_variables
     @GoogleBaseHook.fallback_to_default_project_id
@@ -1351,6 +1356,132 @@ class DataflowHook(GoogleBaseHook):
 
         return job_controller._check_dataflow_job_state(job)
 
+    @GoogleBaseHook.fallback_to_default_project_id
+    def create_data_pipeline(
+        self,
+        body: dict,
+        project_id: str,
+        location: str = DEFAULT_DATAFLOW_LOCATION,
+    ):
+        """
+        Create a new Dataflow Data Pipelines instance.
+
+        :param body: The request body (contains instance of Pipeline). See:
+            https://cloud.google.com/dataflow/docs/reference/data-pipelines/rest/v1/projects.locations.pipelines/create#request-body
+        :param project_id: The ID of the GCP project that owns the job.
+        :param location: The location to direct the Data Pipelines instance to (for example us-central1).
+
+        Returns the created Data Pipelines instance in JSON representation.
+        """
+        parent = self.build_parent_name(project_id, location)
+        service = self.get_pipelines_conn()
+        request = (
+            service.projects()
+            .locations()
+            .pipelines()
+            .create(
+                parent=parent,
+                body=body,
+            )
+        )
+        response = request.execute(num_retries=self.num_retries)
+        return response
+
+    @GoogleBaseHook.fallback_to_default_project_id
+    def get_data_pipeline(
+        self,
+        pipeline_name: str,
+        project_id: str,
+        location: str = DEFAULT_DATAFLOW_LOCATION,
+    ) -> dict:
+        """
+        Retrieve a new Dataflow Data Pipelines instance.
+
+        :param pipeline_name: The display name of the pipeline. In example
+            projects/PROJECT_ID/locations/LOCATION_ID/pipelines/PIPELINE_ID it would be the PIPELINE_ID.
+        :param project_id: The ID of the GCP project that owns the job.
+        :param location: The location to direct the Data Pipelines instance to (for example us-central1).
+
+        Returns the created Data Pipelines instance in JSON representation.
+        """
+        parent = self.build_parent_name(project_id, location)
+        service = self.get_pipelines_conn()
+        request = (
+            service.projects()
+            .locations()
+            .pipelines()
+            .get(
+                name=f"{parent}/pipelines/{pipeline_name}",
+            )
+        )
+        response = request.execute(num_retries=self.num_retries)
+        return response
+
+    @GoogleBaseHook.fallback_to_default_project_id
+    def run_data_pipeline(
+        self,
+        pipeline_name: str,
+        project_id: str,
+        location: str = DEFAULT_DATAFLOW_LOCATION,
+    ) -> dict:
+        """
+        Run a Dataflow Data Pipeline Instance.
+
+        :param pipeline_name: The display name of the pipeline. In example
+            projects/PROJECT_ID/locations/LOCATION_ID/pipelines/PIPELINE_ID it would be the PIPELINE_ID.
+        :param project_id: The ID of the GCP project that owns the job.
+        :param location: The location to direct the Data Pipelines instance to (for example us-central1).
+
+        Returns the created Job in JSON representation.
+        """
+        parent = self.build_parent_name(project_id, location)
+        service = self.get_pipelines_conn()
+        request = (
+            service.projects()
+            .locations()
+            .pipelines()
+            .run(
+                name=f"{parent}/pipelines/{pipeline_name}",
+                body={},
+            )
+        )
+        response = request.execute(num_retries=self.num_retries)
+        return response
+
+    @GoogleBaseHook.fallback_to_default_project_id
+    def delete_data_pipeline(
+        self,
+        pipeline_name: str,
+        project_id: str,
+        location: str = DEFAULT_DATAFLOW_LOCATION,
+    ) -> dict | None:
+        """
+        Delete a Dataflow Data Pipelines Instance.
+
+        :param pipeline_name: The display name of the pipeline. In example
+            projects/PROJECT_ID/locations/LOCATION_ID/pipelines/PIPELINE_ID it would be the PIPELINE_ID.
+        :param project_id: The ID of the GCP project that owns the job.
+        :param location: The location to direct the Data Pipelines instance to (for example us-central1).
+
+        Returns the created Job in JSON representation.
+        """
+        parent = self.build_parent_name(project_id, location)
+        service = self.get_pipelines_conn()
+        request = (
+            service.projects()
+            .locations()
+            .pipelines()
+            .delete(
+                name=f"{parent}/pipelines/{pipeline_name}",
+            )
+        )
+        response = request.execute(num_retries=self.num_retries)
+        return response
+
+    @staticmethod
+    def build_parent_name(project_id: str, location: str):
+        return f"projects/{project_id}/locations/{location}"
+
 
 class AsyncDataflowHook(GoogleBaseAsyncHook):
     """Async hook class for dataflow service."""

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

@@ -19,40 +19,30 @@
 
 from __future__ import annotations
 
-from typing import Sequence
+from typing import TYPE_CHECKING
 
-from googleapiclient.discovery import build
+from deprecated import deprecated
 
-from airflow.providers.google.common.hooks.base_google import (
-    GoogleBaseHook,
-)
-
-DEFAULT_DATAPIPELINE_LOCATION = "us-central1"
+from airflow.exceptions import AirflowProviderDeprecationWarning
+from airflow.providers.google.cloud.hooks.dataflow import DataflowHook
+from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
 
+if TYPE_CHECKING:
+    from googleapiclient.discovery import build
 
-class DataPipelineHook(GoogleBaseHook):
-    """
-    Hook for Google Data Pipelines.
+DEFAULT_DATAPIPELINE_LOCATION = "us-central1"
 
-    All the methods in the hook where project_id is used must be called with
-    keyword arguments rather than positional.
-    """
 
-    def __init__(
-        self,
-        gcp_conn_id: str = "google_cloud_default",
-        impersonation_chain: str | Sequence[str] | None = None,
-        **kwargs,
-    ) -> None:
-        super().__init__(
-            gcp_conn_id=gcp_conn_id,
-            impersonation_chain=impersonation_chain,
-        )
+@deprecated(
+    reason="This hook is deprecated and will be removed after 01.12.2024. Please use `DataflowHook`.",
+    category=AirflowProviderDeprecationWarning,
+)
+class DataPipelineHook(DataflowHook):
+    """Hook for Google Data Pipelines."""
 
     def get_conn(self) -> build:
         """Return a Google Cloud Data Pipelines service object."""
-        http_authorized = self._authorize()
-        return build("datapipelines", "v1", http=http_authorized, cache_discovery=False)
+        return super().get_pipelines_conn()
 
     @GoogleBaseHook.fallback_to_default_project_id
     def create_data_pipeline(
@@ -60,31 +50,9 @@ class DataPipelineHook(GoogleBaseHook):
         body: dict,
         project_id: str,
         location: str = DEFAULT_DATAPIPELINE_LOCATION,
-    ) -> None:
-        """
-        Create a new Data Pipelines instance from the Data Pipelines API.
-
-        :param body: The request body (contains instance of Pipeline). See:
-            https://cloud.google.com/dataflow/docs/reference/data-pipelines/rest/v1/projects.locations.pipelines/create#request-body
-        :param project_id: The ID of the GCP project that owns the job.
-        :param location: The location to direct the Data Pipelines instance to (for example us-central1).
-
-        Returns the created Data Pipelines instance in JSON representation.
-        """
-        parent = self.build_parent_name(project_id, location)
-        service = self.get_conn()
-        self.log.info(dir(service.projects().locations()))
-        request = (
-            service.projects()
-            .locations()
-            .pipelines()
-            .create(
-                parent=parent,
-                body=body,
-            )
-        )
-        response = request.execute(num_retries=self.num_retries)
-        return response
+    ) -> dict:
+        """Create a new Data Pipelines instance from the Data Pipelines API."""
+        return super().create_data_pipeline(body=body, project_id=project_id, location=location)
 
     @GoogleBaseHook.fallback_to_default_project_id
     def run_data_pipeline(
@@ -92,30 +60,11 @@ class DataPipelineHook(GoogleBaseHook):
         data_pipeline_name: str,
         project_id: str,
         location: str = DEFAULT_DATAPIPELINE_LOCATION,
-    ) -> None:
-        """
-        Run a Data Pipelines Instance using the Data Pipelines API.
-
-        :param data_pipeline_name:  The display name of the pipeline. In example
-            projects/PROJECT_ID/locations/LOCATION_ID/pipelines/PIPELINE_ID it would be the PIPELINE_ID.
-        :param project_id: The ID of the GCP project that owns the job.
-        :param location: The location to direct the Data Pipelines instance to (for example us-central1).
-
-        Returns the created Job in JSON representation.
-        """
-        parent = self.build_parent_name(project_id, location)
-        service = self.get_conn()
-        request = (
-            service.projects()
-            .locations()
-            .pipelines()
-            .run(
-                name=f"{parent}/pipelines/{data_pipeline_name}",
-                body={},
-            )
+    ) -> dict:
+        """Run a Data Pipelines Instance using the Data Pipelines API."""
+        return super().run_data_pipeline(
+            pipeline_name=data_pipeline_name, project_id=project_id, location=location
         )
-        response = request.execute(num_retries=self.num_retries)
-        return response
 
     @staticmethod
     def build_parent_name(project_id: str, location: str):

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

@@ -1010,6 +1010,27 @@ class GCSHook(GoogleBaseHook):
         self.log.info("The md5Hash of %s is %s", object_name, blob_md5hash)
         return blob_md5hash
 
+    def get_metadata(self, bucket_name: str, object_name: str) -> dict | None:
+        """
+        Get the metadata of an object in Google Cloud Storage.
+
+        :param bucket_name: Name of the Google Cloud Storage bucket where the object is.
+        :param object_name: The name of the object containing the desired metadata
+        :return: The metadata associated with the object
+        """
+        self.log.info("Retrieving the metadata dict of object (%s) in bucket (%s)", object_name, bucket_name)
+        client = self.get_conn()
+        bucket = client.bucket(bucket_name)
+        blob = bucket.get_blob(blob_name=object_name)
+        if blob is None:
+            raise ValueError("Object (%s) not found in bucket (%s)", object_name, bucket_name)
+        blob_metadata = blob.metadata
+        if blob_metadata:
+            self.log.info("Retrieved metadata of object (%s) with %s fields", object_name, len(blob_metadata))
+        else:
+            self.log.info("Metadata of object (%s) is empty or it does not exist", object_name)
+        return blob_metadata
+
     @GoogleBaseHook.fallback_to_default_project_id
     def create_bucket(
         self,

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

@@ -36,6 +36,7 @@ from google.api_core.exceptions import AlreadyExists, GoogleAPICallError
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
 from google.cloud.exceptions import NotFound
 from google.cloud.pubsub_v1 import PublisherClient, SubscriberClient
+from google.cloud.pubsub_v1.types import PublisherOptions
 from google.pubsub_v1.services.subscriber.async_client import SubscriberAsyncClient
 from googleapiclient.errors import HttpError
 
@@ -79,6 +80,7 @@ class PubSubHook(GoogleBaseHook):
         self,
         gcp_conn_id: str = "google_cloud_default",
         impersonation_chain: str | Sequence[str] | None = None,
+        enable_message_ordering: bool = False,
         **kwargs,
     ) -> None:
         if kwargs.get("delegate_to") is not None:
@@ -90,6 +92,7 @@ class PubSubHook(GoogleBaseHook):
             gcp_conn_id=gcp_conn_id,
             impersonation_chain=impersonation_chain,
         )
+        self.enable_message_ordering = enable_message_ordering
         self._client = None
 
     def get_conn(self) -> PublisherClient:
@@ -99,7 +102,13 @@ class PubSubHook(GoogleBaseHook):
         :return: Google Cloud Pub/Sub client object.
         """
         if not self._client:
-            self._client = PublisherClient(credentials=self.get_credentials(), client_info=CLIENT_INFO)
+            self._client = PublisherClient(
+                credentials=self.get_credentials(),
+                client_info=CLIENT_INFO,
+                publisher_options=PublisherOptions(
+                    enable_message_ordering=self.enable_message_ordering,
+                ),
+            )
         return self._client
 
     @cached_property

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

@@ -551,6 +551,8 @@ class AutoMLHook(GoogleBaseHook):
         is_default_version: bool | None = None,
         model_version_aliases: list[str] | None = None,
         model_version_description: str | None = None,
+        window_stride_length: int | None = None,
+        window_max_count: int | None = None,
     ) -> tuple[models.Model | None, str]:
         """
         Create an AutoML Forecasting Training Job.
@@ -703,6 +705,10 @@ class AutoMLHook(GoogleBaseHook):
         :param sync: Whether to execute this method 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 window_stride_length: Optional. Step length used to generate input examples. Every
+            ``window_stride_length`` rows will be used to generate a sliding window.
+        :param window_max_count: Optional. Number of rows that should be used to generate input examples. If the
+            total row count is larger than this number, the input data will be randomly sampled to hit the count.
         """
         if column_transformations:
             warnings.warn(
@@ -758,6 +764,8 @@ class AutoMLHook(GoogleBaseHook):
             is_default_version=is_default_version,
             model_version_aliases=model_version_aliases,
             model_version_description=model_version_description,
+            window_stride_length=window_stride_length,
+            window_max_count=window_max_count,
         )
         training_id = self.extract_training_id(self._job.resource_name)
         if model:

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

@@ -141,6 +141,8 @@ class GenerativeModelHook(GoogleBaseHook):
         self,
         prompt: str,
         location: str,
+        generation_config: dict | None = None,
+        safety_settings: dict | None = None,
         pretrained_model: str = "gemini-pro",
         project_id: str = PROVIDE_PROJECT_ID,
     ) -> str:
@@ -149,17 +151,21 @@ class GenerativeModelHook(GoogleBaseHook):
 
         :param prompt: Required. Inputs or queries that a user or a program gives
             to the Multi-modal model, in order to elicit a specific response.
+        :param location: Required. The ID of the Google Cloud location that the service belongs to.
+        :param generation_config: Optional. Generation configuration settings.
+        :param safety_settings: Optional. Per request settings for blocking unsafe content.
         :param pretrained_model: By default uses the pre-trained model `gemini-pro`,
             supporting prompts with text-only input, including natural language
             tasks, multi-turn text and code chat, and code generation. It can
             output text and code.
-        :param location: Required. The ID of the Google Cloud location that the service belongs to.
         :param project_id: Required. The ID of the Google Cloud project that the service belongs to.
         """
         vertexai.init(project=project_id, location=location, credentials=self.get_credentials())
 
         model = self.get_generative_model(pretrained_model)
-        response = model.generate_content(prompt)
+        response = model.generate_content(
+            contents=[prompt], generation_config=generation_config, safety_settings=safety_settings
+        )
 
         return response.text
 
@@ -170,6 +176,8 @@ class GenerativeModelHook(GoogleBaseHook):
         location: str,
         media_gcs_path: str,
         mime_type: str,
+        generation_config: dict | None = None,
+        safety_settings: dict | None = None,
         pretrained_model: str = "gemini-pro-vision",
         project_id: str = PROVIDE_PROJECT_ID,
     ) -> str:
@@ -178,6 +186,8 @@ class GenerativeModelHook(GoogleBaseHook):
 
         :param prompt: Required. Inputs or queries that a user or a program gives
             to the Multi-modal model, in order to elicit a specific response.
+        :param generation_config: Optional. Generation configuration settings.
+        :param safety_settings: Optional. Per request settings for blocking unsafe content.
         :param pretrained_model: By default uses the pre-trained model `gemini-pro-vision`,
             supporting prompts with text-only input, including natural language
             tasks, multi-turn text and code chat, and code generation. It can
@@ -192,6 +202,8 @@ class GenerativeModelHook(GoogleBaseHook):
 
         model = self.get_generative_model(pretrained_model)
         part = self.get_generative_model_part(media_gcs_path, mime_type)
-        response = model.generate_content([prompt, part])
+        response = model.generate_content(
+            contents=[prompt, part], generation_config=generation_config, safety_settings=safety_settings
+        )
 
         return response.text

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

@@ -126,7 +126,7 @@ class HyperparameterTuningJobHook(GoogleBaseHook):
             base_output_dir=base_output_dir,
             project=project,
             location=location,
-            credentials=self.get_credentials,
+            credentials=self.get_credentials(),
             labels=labels,
             encryption_spec_key_name=encryption_spec_key_name,
             staging_bucket=staging_bucket,

airflow/providers/google/cloud/links/dataflow.py

@@ -30,6 +30,9 @@ if TYPE_CHECKING:
 DATAFLOW_BASE_LINK = "/dataflow/jobs"
 DATAFLOW_JOB_LINK = DATAFLOW_BASE_LINK + "/{region}/{job_id}?project={project_id}"
 
+DATAFLOW_PIPELINE_BASE_LINK = "/dataflow/pipelines"
+DATAFLOW_PIPELINE_LINK = DATAFLOW_PIPELINE_BASE_LINK + "/{location}/{pipeline_name}?project={project_id}"
+
 
 class DataflowJobLink(BaseGoogleLink):
     """Helper class for constructing Dataflow Job Link."""
@@ -51,3 +54,25 @@ class DataflowJobLink(BaseGoogleLink):
             key=DataflowJobLink.key,
             value={"project_id": project_id, "region": region, "job_id": job_id},
         )
+
+
+class DataflowPipelineLink(BaseGoogleLink):
+    """Helper class for constructing Dataflow Pipeline Link."""
+
+    name = "Dataflow Pipeline"
+    key = "dataflow_pipeline_config"
+    format_str = DATAFLOW_PIPELINE_LINK
+
+    @staticmethod
+    def persist(
+        operator_instance: BaseOperator,
+        context: Context,
+        project_id: str | None,
+        location: str | None,
+        pipeline_name: str | None,
+    ):
+        operator_instance.xcom_push(
+            context,
+            key=DataflowPipelineLink.key,
+            value={"project_id": project_id, "location": location, "pipeline_name": pipeline_name},
+        )