apache-airflow-providers-google 10.24.0rc1__tar.gz → 10.25.0__tar.gz

{apache_airflow_providers_google-10.24.0rc1 → apache_airflow_providers_google-10.25.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: apache-airflow-providers-google
-Version: 10.24.0rc1
+Version: 10.25.0
 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>
@@ -21,9 +21,9 @@ Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Classifier: Topic :: System :: Monitoring
 Requires-Dist: PyOpenSSL>=23.0.0
-Requires-Dist: apache-airflow-providers-common-compat>=1.1.0rc0
-Requires-Dist: apache-airflow-providers-common-sql>=1.7.2rc0
-Requires-Dist: apache-airflow>=2.8.0rc0
+Requires-Dist: apache-airflow-providers-common-compat>=1.2.1
+Requires-Dist: apache-airflow-providers-common-sql>=1.7.2
+Requires-Dist: apache-airflow>=2.8.0
 Requires-Dist: asgiref>=3.5.2
 Requires-Dist: dill>=0.2.3
 Requires-Dist: gcloud-aio-auth>=5.2.0
@@ -36,7 +36,7 @@ Requires-Dist: google-api-core>=2.11.0,!=2.16.0,!=2.18.0
 Requires-Dist: google-api-python-client>=2.0.2
 Requires-Dist: google-auth-httplib2>=0.0.1
 Requires-Dist: google-auth>=2.29.0
-Requires-Dist: google-cloud-aiplatform>=1.63.0
+Requires-Dist: google-cloud-aiplatform>=1.70.0
 Requires-Dist: google-cloud-automl>=2.12.0
 Requires-Dist: google-cloud-batch>=0.13.0
 Requires-Dist: google-cloud-bigquery-datatransfer>=3.13.0
@@ -77,29 +77,30 @@ Requires-Dist: grpcio-gcp>=0.2.2
 Requires-Dist: httpx>=0.25.0
 Requires-Dist: immutabledict>=4.2.0
 Requires-Dist: json-merge-patch>=0.2
-Requires-Dist: looker-sdk>=22.4.0
+Requires-Dist: looker-sdk>=22.4.0,!=24.18.0
 Requires-Dist: pandas-gbq>=0.7.0
 Requires-Dist: pandas>=1.5.3,<2.2;python_version<"3.9"
 Requires-Dist: pandas>=2.1.2,<2.2;python_version>="3.9"
 Requires-Dist: proto-plus>=1.19.6
+Requires-Dist: pyarrow>=14.0.1
 Requires-Dist: python-slugify>=7.0.0
 Requires-Dist: sqlalchemy-bigquery>=1.2.1
 Requires-Dist: sqlalchemy-spanner>=1.6.2
 Requires-Dist: tenacity>=8.1.0
-Requires-Dist: apache-airflow-providers-amazon>=2.6.0rc0 ; extra == "amazon"
+Requires-Dist: apache-airflow-providers-amazon>=2.6.0 ; extra == "amazon"
 Requires-Dist: apache-airflow-providers-apache-beam ; extra == "apache.beam"
 Requires-Dist: apache-beam[gcp] ; extra == "apache.beam"
 Requires-Dist: apache-airflow-providers-apache-cassandra ; extra == "apache.cassandra"
-Requires-Dist: apache-airflow-providers-cncf-kubernetes>=7.2.0rc0 ; extra == "cncf.kubernetes"
+Requires-Dist: apache-airflow-providers-cncf-kubernetes>=7.2.0 ; extra == "cncf.kubernetes"
 Requires-Dist: apache-airflow-providers-common-compat ; extra == "common.compat"
 Requires-Dist: apache-airflow-providers-common-sql ; extra == "common.sql"
-Requires-Dist: apache-airflow-providers-facebook>=2.2.0rc0 ; extra == "facebook"
-Requires-Dist: plyvel ; extra == "leveldb"
+Requires-Dist: apache-airflow-providers-facebook>=2.2.0 ; extra == "facebook"
+Requires-Dist: plyvel>=1.5.1 ; extra == "leveldb"
 Requires-Dist: apache-airflow-providers-microsoft-azure ; extra == "microsoft.azure"
 Requires-Dist: apache-airflow-providers-microsoft-mssql ; extra == "microsoft.mssql"
 Requires-Dist: apache-airflow-providers-mysql ; extra == "mysql"
 Requires-Dist: apache-airflow-providers-openlineage ; extra == "openlineage"
-Requires-Dist: apache-airflow-providers-oracle>=3.1.0rc0 ; extra == "oracle"
+Requires-Dist: apache-airflow-providers-oracle>=3.1.0 ; extra == "oracle"
 Requires-Dist: apache-airflow-providers-postgres ; extra == "postgres"
 Requires-Dist: apache-airflow-providers-presto ; extra == "presto"
 Requires-Dist: apache-airflow-providers-salesforce ; extra == "salesforce"
@@ -107,8 +108,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.24.0/changelog.html
-Project-URL: Documentation, https://airflow.apache.org/docs/apache-airflow-providers-google/10.24.0
+Project-URL: Changelog, https://airflow.apache.org/docs/apache-airflow-providers-google/10.25.0/changelog.html
+Project-URL: Documentation, https://airflow.apache.org/docs/apache-airflow-providers-google/10.25.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
@@ -177,7 +178,7 @@ Provides-Extra: trino
 
 Package ``apache-airflow-providers-google``
 
-Release: ``10.24.0.rc1``
+Release: ``10.25.0``
 
 
 Google services including:
@@ -197,7 +198,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.24.0/>`_.
+in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-google/10.25.0/>`_.
 
 Installation
 ------------
@@ -215,7 +216,7 @@ Requirements
 PIP package                                 Version required
 ==========================================  =========================================
 ``apache-airflow``                          ``>=2.8.0``
-``apache-airflow-providers-common-compat``  ``>=1.1.0``
+``apache-airflow-providers-common-compat``  ``>=1.2.1``
 ``apache-airflow-providers-common-sql``     ``>=1.7.2``
 ``asgiref``                                 ``>=3.5.2``
 ``dill``                                    ``>=0.2.3``
@@ -229,7 +230,7 @@ PIP package                                 Version required
 ``google-api-python-client``                ``>=2.0.2``
 ``google-auth``                             ``>=2.29.0``
 ``google-auth-httplib2``                    ``>=0.0.1``
-``google-cloud-aiplatform``                 ``>=1.63.0``
+``google-cloud-aiplatform``                 ``>=1.70.0``
 ``google-cloud-automl``                     ``>=2.12.0``
 ``google-cloud-bigquery``                   ``!=3.21.*,!=3.22.0,!=3.23.*,>=3.4.0``
 ``google-cloud-bigquery-datatransfer``      ``>=3.13.0``
@@ -269,11 +270,12 @@ PIP package                                 Version required
 ``grpcio-gcp``                              ``>=0.2.2``
 ``httpx``                                   ``>=0.25.0``
 ``json-merge-patch``                        ``>=0.2``
-``looker-sdk``                              ``>=22.4.0``
+``looker-sdk``                              ``>=22.4.0,!=24.18.0``
 ``pandas-gbq``                              ``>=0.7.0``
 ``pandas``                                  ``>=2.1.2,<2.2; python_version >= "3.9"``
 ``pandas``                                  ``>=1.5.3,<2.2; python_version < "3.9"``
 ``proto-plus``                              ``>=1.19.6``
+``pyarrow``                                 ``>=14.0.1``
 ``python-slugify``                          ``>=7.0.0``
 ``PyOpenSSL``                               ``>=23.0.0``
 ``sqlalchemy-bigquery``                     ``>=1.2.1``
@@ -319,4 +321,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.24.0/changelog.html>`_.
+`changelog <https://airflow.apache.org/docs/apache-airflow-providers-google/10.25.0/changelog.html>`_.

{apache_airflow_providers_google-10.24.0rc1 → apache_airflow_providers_google-10.25.0}/README.rst

@@ -42,7 +42,7 @@
 
 Package ``apache-airflow-providers-google``
 
-Release: ``10.24.0.rc1``
+Release: ``10.25.0``
 
 
 Google services including:
@@ -62,7 +62,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.24.0/>`_.
+in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-google/10.25.0/>`_.
 
 Installation
 ------------
@@ -80,7 +80,7 @@ Requirements
 PIP package                                 Version required
 ==========================================  =========================================
 ``apache-airflow``                          ``>=2.8.0``
-``apache-airflow-providers-common-compat``  ``>=1.1.0``
+``apache-airflow-providers-common-compat``  ``>=1.2.1``
 ``apache-airflow-providers-common-sql``     ``>=1.7.2``
 ``asgiref``                                 ``>=3.5.2``
 ``dill``                                    ``>=0.2.3``
@@ -94,7 +94,7 @@ PIP package                                 Version required
 ``google-api-python-client``                ``>=2.0.2``
 ``google-auth``                             ``>=2.29.0``
 ``google-auth-httplib2``                    ``>=0.0.1``
-``google-cloud-aiplatform``                 ``>=1.63.0``
+``google-cloud-aiplatform``                 ``>=1.70.0``
 ``google-cloud-automl``                     ``>=2.12.0``
 ``google-cloud-bigquery``                   ``!=3.21.*,!=3.22.0,!=3.23.*,>=3.4.0``
 ``google-cloud-bigquery-datatransfer``      ``>=3.13.0``
@@ -134,11 +134,12 @@ PIP package                                 Version required
 ``grpcio-gcp``                              ``>=0.2.2``
 ``httpx``                                   ``>=0.25.0``
 ``json-merge-patch``                        ``>=0.2``
-``looker-sdk``                              ``>=22.4.0``
+``looker-sdk``                              ``>=22.4.0,!=24.18.0``
 ``pandas-gbq``                              ``>=0.7.0``
 ``pandas``                                  ``>=2.1.2,<2.2; python_version >= "3.9"``
 ``pandas``                                  ``>=1.5.3,<2.2; python_version < "3.9"``
 ``proto-plus``                              ``>=1.19.6``
+``pyarrow``                                 ``>=14.0.1``
 ``python-slugify``                          ``>=7.0.0``
 ``PyOpenSSL``                               ``>=23.0.0``
 ``sqlalchemy-bigquery``                     ``>=1.2.1``
@@ -184,4 +185,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.24.0/changelog.html>`_.
+`changelog <https://airflow.apache.org/docs/apache-airflow-providers-google/10.25.0/changelog.html>`_.

{apache_airflow_providers_google-10.24.0rc1 → apache_airflow_providers_google-10.25.0}/airflow/providers/google/__init__.py

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

apache_airflow_providers_google-10.25.0/airflow/providers/google/assets/gcs.py

@@ -0,0 +1,45 @@
+# 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 TYPE_CHECKING
+
+from airflow.providers.common.compat.assets import Asset
+from airflow.providers.google.cloud.hooks.gcs import _parse_gcs_url
+
+if TYPE_CHECKING:
+    from urllib.parse import SplitResult
+
+    from airflow.providers.common.compat.openlineage.facet import Dataset as OpenLineageDataset
+
+
+def create_asset(*, bucket: str, key: str, extra: dict | None = None) -> Asset:
+    return Asset(uri=f"gs://{bucket}/{key}", extra=extra)
+
+
+def sanitize_uri(uri: SplitResult) -> SplitResult:
+    if not uri.netloc:
+        raise ValueError("URI format gs:// must contain a bucket name")
+    return uri
+
+
+def convert_asset_to_openlineage(asset: Asset, lineage_context) -> OpenLineageDataset:
+    """Translate Asset with valid AIP-60 uri to OpenLineage with assistance from the hook."""
+    from airflow.providers.common.compat.openlineage.facet import Dataset as OpenLineageDataset
+
+    bucket, key = _parse_gcs_url(asset.uri)
+    return OpenLineageDataset(namespace=f"gs://{bucket}", name=key if key else "/")

{apache_airflow_providers_google-10.24.0rc1 → apache_airflow_providers_google-10.25.0}/airflow/providers/google/cloud/hooks/dataproc.py

@@ -999,6 +999,7 @@ class DataprocHook(GoogleBaseHook):
             individual attempt.
         :param metadata: Additional metadata that is provided to the method.
         """
+        self.log.debug("Creating batch: %s", batch)
         client = self.get_batch_client(region)
         parent = f"projects/{project_id}/regions/{region}"
 

{apache_airflow_providers_google-10.24.0rc1 → apache_airflow_providers_google-10.25.0}/airflow/providers/google/cloud/hooks/gcs.py

@@ -43,6 +43,7 @@ from google.cloud.storage.retry import DEFAULT_RETRY
 from requests import Session
 
 from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
+from airflow.providers.common.compat.lineage.hook import get_hook_lineage_collector
 from airflow.providers.google.cloud.utils.helpers import normalize_directory_path
 from airflow.providers.google.common.consts import CLIENT_INFO
 from airflow.providers.google.common.hooks.base_google import (
@@ -214,6 +215,16 @@ class GCSHook(GoogleBaseHook):
         destination_object = source_bucket.copy_blob(  # type: ignore[attr-defined]
             blob=source_object, destination_bucket=destination_bucket, new_name=destination_object
         )
+        get_hook_lineage_collector().add_input_asset(
+            context=self,
+            scheme="gs",
+            asset_kwargs={"bucket": source_bucket.name, "key": source_object.name},  # type: ignore[attr-defined]
+        )
+        get_hook_lineage_collector().add_output_asset(
+            context=self,
+            scheme="gs",
+            asset_kwargs={"bucket": destination_bucket.name, "key": destination_object.name},  # type: ignore[union-attr]
+        )
 
         self.log.info(
             "Object %s in bucket %s copied to object %s in bucket %s",
@@ -267,6 +278,16 @@ class GCSHook(GoogleBaseHook):
             ).rewrite(source=source_object, token=token)
 
             self.log.info("Total Bytes: %s | Bytes Written: %s", total_bytes, bytes_rewritten)
+        get_hook_lineage_collector().add_input_asset(
+            context=self,
+            scheme="gs",
+            asset_kwargs={"bucket": source_bucket.name, "key": source_object.name},  # type: ignore[attr-defined]
+        )
+        get_hook_lineage_collector().add_output_asset(
+            context=self,
+            scheme="gs",
+            asset_kwargs={"bucket": destination_bucket.name, "key": destination_object},  # type: ignore[attr-defined]
+        )
         self.log.info(
             "Object %s in bucket %s rewritten to object %s in bucket %s",
             source_object.name,  # type: ignore[attr-defined]
@@ -345,9 +366,18 @@ class GCSHook(GoogleBaseHook):
 
                 if filename:
                     blob.download_to_filename(filename, timeout=timeout)
+                    get_hook_lineage_collector().add_input_asset(
+                        context=self, scheme="gs", asset_kwargs={"bucket": bucket.name, "key": blob.name}
+                    )
+                    get_hook_lineage_collector().add_output_asset(
+                        context=self, scheme="file", asset_kwargs={"path": filename}
+                    )
                     self.log.info("File downloaded to %s", filename)
                     return filename
                 else:
+                    get_hook_lineage_collector().add_input_asset(
+                        context=self, scheme="gs", asset_kwargs={"bucket": bucket.name, "key": blob.name}
+                    )
                     return blob.download_as_bytes()
 
             except GoogleCloudError:
@@ -555,6 +585,9 @@ class GCSHook(GoogleBaseHook):
             _call_with_retry(
                 partial(blob.upload_from_filename, filename=filename, content_type=mime_type, timeout=timeout)
             )
+            get_hook_lineage_collector().add_input_asset(
+                context=self, scheme="file", asset_kwargs={"path": filename}
+            )
 
             if gzip:
                 os.remove(filename)
@@ -576,6 +609,10 @@ class GCSHook(GoogleBaseHook):
         else:
             raise ValueError("'filename' and 'data' parameter missing. One is required to upload to gcs.")
 
+        get_hook_lineage_collector().add_output_asset(
+            context=self, scheme="gs", asset_kwargs={"bucket": bucket.name, "key": blob.name}
+        )
+
     def exists(self, bucket_name: str, object_name: str, retry: Retry = DEFAULT_RETRY) -> bool:
         """
         Check for the existence of a file in Google Cloud Storage.
@@ -691,6 +728,9 @@ class GCSHook(GoogleBaseHook):
         bucket = client.bucket(bucket_name)
         blob = bucket.blob(blob_name=object_name)
         blob.delete()
+        get_hook_lineage_collector().add_input_asset(
+            context=self, scheme="gs", asset_kwargs={"bucket": bucket.name, "key": blob.name}
+        )
 
         self.log.info("Blob %s deleted.", object_name)
 
@@ -1198,9 +1238,17 @@ class GCSHook(GoogleBaseHook):
         client = self.get_conn()
         bucket = client.bucket(bucket_name)
         destination_blob = bucket.blob(destination_object)
-        destination_blob.compose(
-            sources=[bucket.blob(blob_name=source_object) for source_object in source_objects]
+        source_blobs = [bucket.blob(blob_name=source_object) for source_object in source_objects]
+        destination_blob.compose(sources=source_blobs)
+        get_hook_lineage_collector().add_output_asset(
+            context=self, scheme="gs", asset_kwargs={"bucket": bucket.name, "key": destination_blob.name}
         )
+        for single_source_blob in source_blobs:
+            get_hook_lineage_collector().add_input_asset(
+                context=self,
+                scheme="gs",
+                asset_kwargs={"bucket": bucket.name, "key": single_source_blob.name},
+            )
 
         self.log.info("Completed successfully.")
 

{apache_airflow_providers_google-10.24.0rc1 → apache_airflow_providers_google-10.25.0}/airflow/providers/google/cloud/hooks/vertex_ai/generative_model.py

@@ -20,12 +20,15 @@
 from __future__ import annotations
 
 import time
+from datetime import timedelta
 from typing import TYPE_CHECKING, Sequence
 
 import vertexai
 from vertexai.generative_models import GenerativeModel, Part
 from vertexai.language_models import TextEmbeddingModel, TextGenerationModel
+from vertexai.preview.caching import CachedContent
 from vertexai.preview.evaluation import EvalResult, EvalTask
+from vertexai.preview.generative_models import GenerativeModel as preview_generative_model
 from vertexai.preview.tuning import sft
 
 from airflow.exceptions import AirflowProviderDeprecationWarning
@@ -95,6 +98,16 @@ class GenerativeModelHook(GoogleBaseHook):
         )
         return eval_task
 
+    def get_cached_context_model(
+        self,
+        cached_content_name: str,
+    ) -> preview_generative_model:
+        """Return a Generative Model with Cached Context."""
+        cached_content = CachedContent(cached_content_name=cached_content_name)
+
+        cached_context_model = preview_generative_model.from_cached_content(cached_content)
+        return cached_context_model
+
     @deprecated(
         planned_removal_date="January 01, 2025",
         use_instead="Part objects included in contents parameter of "
@@ -528,3 +541,69 @@ class GenerativeModelHook(GoogleBaseHook):
         )
 
         return eval_result
+
+    def create_cached_content(
+        self,
+        model_name: str,
+        location: str,
+        ttl_hours: float = 1,
+        system_instruction: str | None = None,
+        contents: list | None = None,
+        display_name: str | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
+    ) -> str:
+        """
+        Create CachedContent to reduce the cost of requests that contain repeat content with high input token counts.
+
+        :param project_id: Required. The ID of the Google Cloud project that the service belongs to.
+        :param location: Required. The ID of the Google Cloud location that the service belongs to.
+        :param model_name: Required. The name of the publisher model to use for cached content.
+        :param system_instruction: Developer set system instruction.
+        :param contents: The content to cache.
+        :param ttl_hours: The TTL for this resource in hours. The expiration time is computed: now + TTL.
+            Defaults to one hour.
+        :param display_name: The user-generated meaningful display name of the cached content
+        """
+        vertexai.init(project=project_id, location=location, credentials=self.get_credentials())
+
+        response = CachedContent.create(
+            model_name=model_name,
+            system_instruction=system_instruction,
+            contents=contents,
+            ttl=timedelta(hours=ttl_hours),
+            display_name=display_name,
+        )
+
+        return response.name
+
+    def generate_from_cached_content(
+        self,
+        location: str,
+        cached_content_name: str,
+        contents: list,
+        generation_config: dict | None = None,
+        safety_settings: dict | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
+    ) -> str:
+        """
+        Generate a response from CachedContent.
+
+        :param project_id: Required. The ID of the Google Cloud project that the service belongs to.
+        :param location: Required. The ID of the Google Cloud location that the service belongs to.
+        :param cached_content_name: Required. The name of the cached content resource.
+        :param contents: Required. The multi-part content of a message that a user or a program
+            gives to the generative 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.
+        """
+        vertexai.init(project=project_id, location=location, credentials=self.get_credentials())
+
+        cached_context_model = self.get_cached_context_model(cached_content_name=cached_content_name)
+
+        response = cached_context_model.generate_content(
+            contents=contents,
+            generation_config=generation_config,
+            safety_settings=safety_settings,
+        )
+
+        return response.text

{apache_airflow_providers_google-10.24.0rc1 → apache_airflow_providers_google-10.25.0}/airflow/providers/google/cloud/links/translate.py

@@ -167,13 +167,14 @@ class TranslationLegacyModelPredictLink(BaseGoogleLink):
         task_instance,
         model_id: str,
         project_id: str,
+        dataset_id: str,
     ):
         task_instance.xcom_push(
             context,
             key=TranslationLegacyModelPredictLink.key,
             value={
                 "location": task_instance.location,
-                "dataset_id": task_instance.model.dataset_id,
+                "dataset_id": dataset_id,
                 "model_id": model_id,
                 "project_id": project_id,
             },

{apache_airflow_providers_google-10.24.0rc1 → apache_airflow_providers_google-10.25.0}/airflow/providers/google/cloud/operators/automl.py

@@ -22,7 +22,7 @@ from __future__ import annotations
 import ast
 import warnings
 from functools import cached_property
-from typing import TYPE_CHECKING, Sequence, Tuple
+from typing import TYPE_CHECKING, Sequence, Tuple, cast
 
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
 from google.cloud.automl_v1beta1 import (
@@ -280,17 +280,22 @@ class AutoMLPredictOperator(GoogleCloudBaseOperator):
                 impersonation_chain=self.impersonation_chain,
             )
 
+    @cached_property
+    def model(self) -> Model | None:
+        if self.model_id:
+            hook = cast(CloudAutoMLHook, self.hook)
+            return hook.get_model(
+                model_id=self.model_id,
+                location=self.location,
+                project_id=self.project_id,
+                retry=self.retry,
+                timeout=self.timeout,
+                metadata=self.metadata,
+            )
+        return None
+
     def _check_model_type(self):
-        hook = self.hook
-        model = hook.get_model(
-            model_id=self.model_id,
-            location=self.location,
-            project_id=self.project_id,
-            retry=self.retry,
-            timeout=self.timeout,
-            metadata=self.metadata,
-        )
-        if not hasattr(model, "translation_model_metadata"):
+        if not hasattr(self.model, "translation_model_metadata"):
             raise AirflowException(
                 "AutoMLPredictOperator for text, image, and video prediction has been deprecated. "
                 "Please use endpoint_id param instead of model_id param."
@@ -329,11 +334,13 @@ class AutoMLPredictOperator(GoogleCloudBaseOperator):
             )
 
         project_id = self.project_id or hook.project_id
-        if project_id and self.model_id:
+        dataset_id: str | None = self.model.dataset_id if self.model else None
+        if project_id and self.model_id and dataset_id:
             TranslationLegacyModelPredictLink.persist(
                 context=context,
                 task_instance=self,
                 model_id=self.model_id,
+                dataset_id=dataset_id,
                 project_id=project_id,
             )
         return PredictResponse.to_dict(result)
@@ -431,12 +438,16 @@ class AutoMLBatchPredictOperator(GoogleCloudBaseOperator):
         self.input_config = input_config
         self.output_config = output_config
 
-    def execute(self, context: Context):
-        hook = CloudAutoMLHook(
+    @cached_property
+    def hook(self) -> CloudAutoMLHook:
+        return CloudAutoMLHook(
             gcp_conn_id=self.gcp_conn_id,
             impersonation_chain=self.impersonation_chain,
         )
-        self.model: Model = hook.get_model(
+
+    @cached_property
+    def model(self) -> Model:
+        return self.hook.get_model(
             model_id=self.model_id,
             location=self.location,
             project_id=self.project_id,
@@ -445,6 +456,7 @@ class AutoMLBatchPredictOperator(GoogleCloudBaseOperator):
             metadata=self.metadata,
         )
 
+    def execute(self, context: Context):
         if not hasattr(self.model, "translation_model_metadata"):
             _raise_exception_for_deprecated_operator(
                 self.__class__.__name__,
@@ -456,7 +468,7 @@ class AutoMLBatchPredictOperator(GoogleCloudBaseOperator):
                 ],
             )
         self.log.info("Fetch batch prediction.")
-        operation = hook.batch_predict(
+        operation = self.hook.batch_predict(
             model_id=self.model_id,
             input_config=self.input_config,
             output_config=self.output_config,
@@ -467,16 +479,17 @@ class AutoMLBatchPredictOperator(GoogleCloudBaseOperator):
             timeout=self.timeout,
             metadata=self.metadata,
         )
-        operation_result = hook.wait_for_operation(timeout=self.timeout, operation=operation)
+        operation_result = self.hook.wait_for_operation(timeout=self.timeout, operation=operation)
         result = BatchPredictResult.to_dict(operation_result)
         self.log.info("Batch prediction is ready.")
-        project_id = self.project_id or hook.project_id
+        project_id = self.project_id or self.hook.project_id
         if project_id:
             TranslationLegacyModelPredictLink.persist(
                 context=context,
                 task_instance=self,
                 model_id=self.model_id,
                 project_id=project_id,
+                dataset_id=self.model.dataset_id,
             )
         return result
 

{apache_airflow_providers_google-10.24.0rc1 → apache_airflow_providers_google-10.25.0}/airflow/providers/google/cloud/operators/bigquery.py

@@ -35,7 +35,7 @@ from airflow.configuration import conf
 from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning, AirflowSkipException
 from airflow.models import BaseOperator, BaseOperatorLink
 from airflow.models.xcom import XCom
-from airflow.providers.common.sql.operators.sql import (
+from airflow.providers.common.sql.operators.sql import (  # type: ignore[attr-defined] # for _parse_boolean
     SQLCheckOperator,
     SQLColumnCheckOperator,
     SQLIntervalCheckOperator,

{apache_airflow_providers_google-10.24.0rc1 → apache_airflow_providers_google-10.25.0}/airflow/providers/google/cloud/operators/cloud_run.py

@@ -243,18 +243,16 @@ class CloudRunListJobsOperator(GoogleCloudBaseOperator):
 
 class CloudRunExecuteJobOperator(GoogleCloudBaseOperator):
     """
-    Executes a job and wait for the operation to be completed. Pushes the executed job to xcom.
+    Executes a job and waits for the operation to be completed. Pushes the executed job to xcom.
 
     :param project_id: Required. The ID of the Google Cloud project that the service belongs to.
     :param region: Required. The ID of the Google Cloud region that the service belongs to.
     :param job_name: Required. The name of the job to update.
-    :param job: Required. The job descriptor containing the new configuration of the job to update.
-        The name field will be replaced by job_name
     :param overrides: Optional map of override values.
     :param gcp_conn_id: The connection ID used to connect to Google Cloud.
-    :param polling_period_seconds: Optional: Control the rate of the poll for the result of deferrable run.
+    :param polling_period_seconds: Optional. Control the rate of the poll for the result of deferrable run.
         By default, the trigger will poll every 10 seconds.
-    :param timeout: The timeout for this request.
+    :param timeout_seconds: Optional. The timeout for this request, in seconds.
     :param impersonation_chain: Optional service account to impersonate using short-term
         credentials, or chained list of accounts required to get the access_token
         of the last account in the list, which will be impersonated in the request.
@@ -263,7 +261,7 @@ class CloudRunExecuteJobOperator(GoogleCloudBaseOperator):
         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 deferrable: Run operator in the deferrable mode
+    :param deferrable: Run the operator in deferrable mode.
     """
 
     template_fields = ("project_id", "region", "gcp_conn_id", "impersonation_chain", "job_name", "overrides")