apache-airflow-providers-google 10.20.0rc1__py3-none-any.whl → 10.21.0rc1__py3-none-any.whl

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

@@ -104,7 +104,8 @@ class GKEClusterConnection:
 
 
 class GKEHook(GoogleBaseHook):
-    """Google Kubernetes Engine cluster APIs.
+    """
+    Google Kubernetes Engine cluster APIs.
 
     All the methods in the hook where project_id is used must be called with
     keyword arguments rather than positional.
@@ -157,7 +158,8 @@ class GKEHook(GoogleBaseHook):
         return self.get_conn()
 
     def wait_for_operation(self, operation: Operation, project_id: str = PROVIDE_PROJECT_ID) -> Operation:
-        """Continuously fetch the status from Google Cloud.
+        """
+        Continuously fetch the status from Google Cloud.
 
         This is done until the given operation completes, or raises an error.
 
@@ -177,7 +179,8 @@ class GKEHook(GoogleBaseHook):
         return operation
 
     def get_operation(self, operation_name: str, project_id: str = PROVIDE_PROJECT_ID) -> Operation:
-        """Get an operation from Google Cloud.
+        """
+        Get an operation from Google Cloud.
 
         :param operation_name: Name of operation to fetch
         :param project_id: Google Cloud project ID
@@ -192,7 +195,8 @@ class GKEHook(GoogleBaseHook):
 
     @staticmethod
     def _append_label(cluster_proto: Cluster, key: str, val: str) -> Cluster:
-        """Append labels to provided Cluster Protobuf.
+        """
+        Append labels to provided Cluster Protobuf.
 
         Labels must fit the regex ``[a-z]([-a-z0-9]*[a-z0-9])?`` (current
          airflow version string follows semantic versioning spec: x.y.z).
@@ -216,7 +220,8 @@ class GKEHook(GoogleBaseHook):
         retry: Retry | _MethodDefault = DEFAULT,
         timeout: float | None = None,
     ) -> Operation | None:
-        """Delete the cluster, the Kubernetes endpoint, and all worker nodes.
+        """
+        Delete the cluster, the Kubernetes endpoint, and all worker nodes.
 
         Firewalls and routes that were configured during cluster creation are
         also deleted. Other Google Compute Engine resources that might be in use
@@ -259,7 +264,8 @@ class GKEHook(GoogleBaseHook):
         retry: Retry | _MethodDefault = DEFAULT,
         timeout: float | None = None,
     ) -> Operation | Cluster:
-        """Create a cluster.
+        """
+        Create a cluster.
 
         This should consist of the specified number, and the type of Google
         Compute Engine instances.
@@ -314,7 +320,8 @@ class GKEHook(GoogleBaseHook):
         retry: Retry | _MethodDefault = DEFAULT,
         timeout: float | None = None,
     ) -> Cluster:
-        """Get details of specified cluster.
+        """
+        Get details of specified cluster.
 
         :param name: The name of the cluster to retrieve.
         :param project_id: Google Cloud project ID.
@@ -404,7 +411,8 @@ class GKEAsyncHook(GoogleBaseAsyncHook):
         operation_name: str,
         project_id: str = PROVIDE_PROJECT_ID,
     ) -> Operation:
-        """Fetch an operation from Google Cloud.
+        """
+        Fetch an operation from Google Cloud.
 
         :param operation_name: Name of operation to fetch.
         :param project_id: Google Cloud project ID.
@@ -420,7 +428,8 @@ class GKEAsyncHook(GoogleBaseAsyncHook):
 
 
 class GKEKubernetesHook(GoogleBaseHook, KubernetesHook):
-    """GKE authenticated hook for standard Kubernetes API.
+    """
+    GKE authenticated hook for standard Kubernetes API.
 
     This hook provides full set of the standard Kubernetes API provided by the KubernetesHook,
     and at the same time it provides a GKE authentication, so it makes it possible to KubernetesHook
@@ -506,7 +515,8 @@ class GKEKubernetesHook(GoogleBaseHook, KubernetesHook):
 
 
 class GKEKubernetesAsyncHook(GoogleBaseAsyncHook, AsyncKubernetesHook):
-    """Async GKE authenticated hook for standard Kubernetes API.
+    """
+    Async GKE authenticated hook for standard Kubernetes API.
 
     This hook provides full set of the standard Kubernetes API provided by the AsyncKubernetesHook,
     and at the same time it provides a GKE authentication, so it makes it possible to KubernetesHook
@@ -639,7 +649,8 @@ class GKEJobHook(GKEKubernetesHook):
     category=AirflowProviderDeprecationWarning,
 )
 class GKEPodAsyncHook(GKEKubernetesAsyncHook):
-    """Google Kubernetes Engine pods APIs asynchronously.
+    """
+    Google Kubernetes Engine pods APIs asynchronously.
 
     :param cluster_url: The URL pointed to the cluster.
     :param ssl_ca_cert: SSL certificate used for authentication to the pod.

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

@@ -14,7 +14,8 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
-"""OS Login hooks.
+"""
+OS Login hooks.
 
 .. spelling:word-list::
     ImportSshPublicKeyResponse

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

@@ -111,21 +111,24 @@ class SecretsManagerHook(GoogleBaseHook):
 
 
 class GoogleCloudSecretManagerHook(GoogleBaseHook):
-    """Hook for the Google Cloud Secret Manager API.
+    """
+    Hook for the Google Cloud Secret Manager API.
 
     See https://cloud.google.com/secret-manager
     """
 
     @cached_property
     def client(self):
-        """Create a Secret Manager Client.
+        """
+        Create a Secret Manager Client.
 
         :return: Secret Manager client.
         """
         return SecretManagerServiceClient(credentials=self.get_credentials(), client_info=CLIENT_INFO)
 
     def get_conn(self) -> SecretManagerServiceClient:
-        """Retrieve the connection to Secret Manager.
+        """
+        Retrieve the connection to Secret Manager.
 
         :return: Secret Manager client.
         """
@@ -141,7 +144,8 @@ class GoogleCloudSecretManagerHook(GoogleBaseHook):
         timeout: float | None = None,
         metadata: Sequence[tuple[str, str]] = (),
     ) -> Secret:
-        """Create a secret.
+        """
+        Create a secret.
 
         .. seealso::
             For more details see API documentation:
@@ -180,7 +184,8 @@ class GoogleCloudSecretManagerHook(GoogleBaseHook):
         timeout: float | None = None,
         metadata: Sequence[tuple[str, str]] = (),
     ) -> SecretVersion:
-        """Add a version to the secret.
+        """
+        Add a version to the secret.
 
         .. seealso::
             For more details see API documentation:
@@ -218,7 +223,8 @@ class GoogleCloudSecretManagerHook(GoogleBaseHook):
         timeout: float | None = None,
         metadata: Sequence[tuple[str, str]] = (),
     ) -> ListSecretsPager:
-        """List secrets.
+        """
+        List secrets.
 
         .. seealso::
             For more details see API documentation:
@@ -250,7 +256,8 @@ class GoogleCloudSecretManagerHook(GoogleBaseHook):
 
     @GoogleBaseHook.fallback_to_default_project_id
     def secret_exists(self, project_id: str, secret_id: str) -> bool:
-        """Check whether secret exists.
+        """
+        Check whether secret exists.
 
         :param project_id: Required. ID of the GCP project that owns the job.
             If set to ``None`` or missing, the default project_id from the GCP connection is used.
@@ -276,7 +283,8 @@ class GoogleCloudSecretManagerHook(GoogleBaseHook):
         timeout: float | None = None,
         metadata: Sequence[tuple[str, str]] = (),
     ) -> AccessSecretVersionResponse:
-        """Access a secret version.
+        """
+        Access a secret version.
 
         .. seealso::
             For more details see API documentation:
@@ -311,7 +319,8 @@ class GoogleCloudSecretManagerHook(GoogleBaseHook):
         timeout: float | None = None,
         metadata: Sequence[tuple[str, str]] = (),
     ) -> None:
-        """Delete a secret.
+        """
+        Delete a secret.
 
         .. seealso::
             For more details see API documentation:

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

@@ -71,7 +71,8 @@ class CloudTranslateHook(GoogleBaseHook):
         source_language: str | None = None,
         model: str | list[str] | None = None,
     ) -> dict:
-        """Translate a string or list of strings.
+        """
+        Translate a string or list of strings.
 
         See https://cloud.google.com/translate/docs/translating-text
 

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

@@ -551,7 +551,8 @@ class BatchPredictionJobAsyncHook(GoogleBaseAsyncHook):
         timeout: float | None = None,
         metadata: Sequence[tuple[str, str]] = (),
     ) -> types.BatchPredictionJob:
-        """Retrieve a batch prediction tuning job.
+        """
+        Retrieve a batch prediction tuning job.
 
         :param project_id: Required. The ID of the Google Cloud project that the job belongs to.
         :param location: Required. The ID of the Google Cloud region that the job belongs to.

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

@@ -22,9 +22,11 @@ from __future__ import annotations
 from typing import Sequence
 
 import vertexai
+from deprecated import deprecated
 from vertexai.generative_models import GenerativeModel, Part
 from vertexai.language_models import TextEmbeddingModel, TextGenerationModel
 
+from airflow.exceptions import AirflowProviderDeprecationWarning
 from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID, GoogleBaseHook
 
 
@@ -59,11 +61,23 @@ class GenerativeModelHook(GoogleBaseHook):
         model = GenerativeModel(pretrained_model)
         return model
 
+    @deprecated(
+        reason=(
+            "The `get_generative_model_part` method is deprecated and will be removed after 01.01.2025, please include `Part` objects in `contents` parameter of `airflow.providers.google.cloud.hooks.generative_model.GenerativeModelHook.generative_model_generate_content`"
+        ),
+        category=AirflowProviderDeprecationWarning,
+    )
     def get_generative_model_part(self, content_gcs_path: str, content_mime_type: str | None = None) -> Part:
         """Return a Generative Model Part object."""
         part = Part.from_uri(content_gcs_path, mime_type=content_mime_type)
         return part
 
+    @deprecated(
+        reason=(
+            "The `prompt_language_model` method is deprecated and will be removed after 01.01.2025, please use `airflow.providers.google.cloud.hooks.generative_model.GenerativeModelHook.text_generation_model_predict` method."
+        ),
+        category=AirflowProviderDeprecationWarning,
+    )
     @GoogleBaseHook.fallback_to_default_project_id
     def prompt_language_model(
         self,
@@ -112,6 +126,12 @@ class GenerativeModelHook(GoogleBaseHook):
         )
         return response.text
 
+    @deprecated(
+        reason=(
+            "The `generate_text_embeddings` method is deprecated and will be removed after 01.01.2025, please use `airflow.providers.google.cloud.hooks.generative_model.GenerativeModelHook.text_embedding_model_get_embeddings` method."
+        ),
+        category=AirflowProviderDeprecationWarning,
+    )
     @GoogleBaseHook.fallback_to_default_project_id
     def generate_text_embeddings(
         self,
@@ -136,6 +156,12 @@ class GenerativeModelHook(GoogleBaseHook):
 
         return response.values
 
+    @deprecated(
+        reason=(
+            "The `prompt_multimodal_model` method is deprecated and will be removed after 01.01.2025, please use `airflow.providers.google.cloud.hooks.generative_model.GenerativeModelHook.generative_model_generate_content` method."
+        ),
+        category=AirflowProviderDeprecationWarning,
+    )
     @GoogleBaseHook.fallback_to_default_project_id
     def prompt_multimodal_model(
         self,
@@ -169,6 +195,12 @@ class GenerativeModelHook(GoogleBaseHook):
 
         return response.text
 
+    @deprecated(
+        reason=(
+            "The `prompt_multimodal_model_with_media` method is deprecated and will be removed after 01.01.2025, please use `airflow.providers.google.cloud.hooks.generative_model.GenerativeModelHook.generative_model_generate_content` method."
+        ),
+        category=AirflowProviderDeprecationWarning,
+    )
     @GoogleBaseHook.fallback_to_default_project_id
     def prompt_multimodal_model_with_media(
         self,
@@ -207,3 +239,112 @@ class GenerativeModelHook(GoogleBaseHook):
         )
 
         return response.text
+
+    @GoogleBaseHook.fallback_to_default_project_id
+    def text_generation_model_predict(
+        self,
+        prompt: str,
+        pretrained_model: str,
+        temperature: float,
+        max_output_tokens: int,
+        top_p: float,
+        top_k: int,
+        location: str,
+        project_id: str = PROVIDE_PROJECT_ID,
+    ) -> str:
+        """
+        Use the Vertex AI PaLM API to generate natural language text.
+
+        :param prompt: Required. Inputs or queries that a user or a program gives
+            to the Vertex AI PaLM API, in order to elicit a specific response.
+        :param pretrained_model: A pre-trained model optimized for performing natural
+            language tasks such as classification, summarization, extraction, content
+            creation, and ideation.
+        :param temperature: Temperature controls the degree of randomness in token
+            selection.
+        :param max_output_tokens: Token limit determines the maximum amount of text
+            output.
+        :param top_p: Tokens are selected from most probable to least until the sum
+            of their probabilities equals the top_p value. Defaults to 0.8.
+        :param top_k: A top_k of 1 means the selected token is the most probable
+            among all tokens.
+        :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())
+
+        parameters = {
+            "temperature": temperature,
+            "max_output_tokens": max_output_tokens,
+            "top_p": top_p,
+            "top_k": top_k,
+        }
+
+        model = self.get_text_generation_model(pretrained_model)
+
+        response = model.predict(
+            prompt=prompt,
+            **parameters,
+        )
+        return response.text
+
+    @GoogleBaseHook.fallback_to_default_project_id
+    def text_embedding_model_get_embeddings(
+        self,
+        prompt: str,
+        pretrained_model: str,
+        location: str,
+        project_id: str = PROVIDE_PROJECT_ID,
+    ) -> list:
+        """
+        Use the Vertex AI PaLM API to generate text embeddings.
+
+        :param prompt: Required. Inputs or queries that a user or a program gives
+            to the Vertex AI PaLM API, in order to elicit a specific response.
+        :param pretrained_model: A pre-trained model optimized for generating text embeddings.
+        :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_text_embedding_model(pretrained_model)
+
+        response = model.get_embeddings([prompt])[0]  # single prompt
+
+        return response.values
+
+    @GoogleBaseHook.fallback_to_default_project_id
+    def generative_model_generate_content(
+        self,
+        contents: list,
+        location: str,
+        tools: list | None = None,
+        generation_config: dict | None = None,
+        safety_settings: dict | None = None,
+        pretrained_model: str = "gemini-pro",
+        project_id: str = PROVIDE_PROJECT_ID,
+    ) -> str:
+        """
+        Use the Vertex AI Gemini Pro foundation model to generate natural language text.
+
+        :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 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 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(
+            contents=contents,
+            tools=tools,
+            generation_config=generation_config,
+            safety_settings=safety_settings,
+        )
+
+        return response.text

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

@@ -15,7 +15,8 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
-"""This module contains a Google Cloud Vertex AI hook.
+"""
+This module contains a Google Cloud Vertex AI hook.
 
 .. spelling:word-list::
 

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

@@ -30,7 +30,8 @@ BASE_LINK = "https://console.cloud.google.com"
 
 
 class BaseGoogleLink(BaseOperatorLink):
-    """Base class for all Google links.
+    """
+    Base class for all Google links.
 
     :meta private:
     """

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

@@ -36,7 +36,8 @@ DATAFUSION_PIPELINE_LINK = "{uri}/pipelines/ns/{namespace}/view/{pipeline_name}"
 
 
 class BaseGoogleLink(BaseOperatorLink):
-    """Link for Google operators.
+    """
+    Link for Google operators.
 
     Prevent adding ``https://console.cloud.google.com`` in front of every link
     where URI is used.

airflow/providers/google/cloud/log/stackdriver_task_handler.py

@@ -50,7 +50,8 @@ _DEFAULT_SCOPESS = frozenset(
 
 
 class StackdriverTaskHandler(logging.Handler):
-    """Handler that directly makes Stackdriver logging API calls.
+    """
+    Handler that directly makes Stackdriver logging API calls.
 
     This is a Python standard ``logging`` handler using that can be used to
     route Python standard logging messages directly to the Stackdriver
@@ -174,7 +175,8 @@ class StackdriverTaskHandler(logging.Handler):
         return labels or {}
 
     def emit(self, record: logging.LogRecord) -> None:
-        """Actually log the specified logging record.
+        """
+        Actually log the specified logging record.
 
         :param record: The record to be logged.
         """

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

@@ -67,8 +67,18 @@ class _BigQueryOpenLineageMixin:
         from airflow.providers.openlineage.sqlparser import SQLParser
 
         if not self.job_id:
+            if hasattr(self, "log"):
+                self.log.warning("No BigQuery job_id was found by OpenLineage.")
             return OperatorLineage()
 
+        if not self.hook:
+            from airflow.providers.google.cloud.hooks.bigquery import BigQueryHook
+
+            self.hook = BigQueryHook(
+                gcp_conn_id=self.gcp_conn_id,
+                impersonation_chain=self.impersonation_chain,
+            )
+
         run_facets: dict[str, BaseFacet] = {
             "externalQuery": ExternalQueryRunFacet(externalQueryId=self.job_id, source="bigquery")
         }

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

@@ -89,7 +89,8 @@ def get_identity_column_lineage_facet(
 
 @define
 class BigQueryJobRunFacet(BaseFacet):
-    """Facet that represents relevant statistics of bigquery run.
+    """
+    Facet that represents relevant statistics of bigquery run.
 
     This facet is used to provide statistics about bigquery run.
 
@@ -134,7 +135,8 @@ class BigQueryErrorRunFacet(BaseFacet):
 
 
 def get_from_nullable_chain(source: Any, chain: list[str]) -> Any | None:
-    """Get object from nested structure of objects, where it's not guaranteed that all keys in the nested structure exist.
+    """
+    Get object from nested structure of objects, where it's not guaranteed that all keys in the nested structure exist.
 
     Intended to replace chain of `dict.get()` statements.