apache-airflow-providers-google 11.0.0__py3-none-any.whl → 12.0.0__py3-none-any.whl

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

@@ -39,9 +39,14 @@ if TYPE_CHECKING:
     from google.api_core.operation import Operation
     from google.cloud.translate_v3.services.translation_service import pagers
     from google.cloud.translate_v3.types import (
+        BatchDocumentInputConfig,
+        BatchDocumentOutputConfig,
         DatasetInputConfig,
+        DocumentInputConfig,
+        DocumentOutputConfig,
         InputConfig,
         OutputConfig,
+        TranslateDocumentResponse,
         TranslateTextGlossaryConfig,
         TransliterationConfig,
         automl_translation,
@@ -560,3 +565,353 @@ class TranslateHook(GoogleBaseHook):
             metadata=metadata,
         )
         return result
+
+    def create_model(
+        self,
+        dataset_id: str,
+        display_name: str,
+        project_id: str,
+        location: str,
+        retry: Retry | _MethodDefault = DEFAULT,
+        timeout: float | None = None,
+        metadata: Sequence[tuple[str, str]] = (),
+    ) -> Operation:
+        """
+        Create the native model by training on translation dataset provided.
+
+        :param dataset_id: ID of dataset to be used for model training.
+        :param display_name: Display name of the model trained.
+            A-Z and a-z, underscores (_), and ASCII digits 0-9.
+        :param project_id: ID of the Google Cloud project where dataset is located. If not provided
+            default project_id is used.
+        :param location: The location of the project.
+        :param retry: A retry object used to retry requests. If `None` is specified, requests will not be
+            retried.
+        :param timeout: The amount of time, in seconds, to wait for the request to complete. Note that if
+            `retry` is specified, the timeout applies to each individual attempt.
+        :param metadata: Additional metadata that is provided to the method.
+
+        :return: `Operation` object with the model creation results, when finished.
+        """
+        client = self.get_client()
+        project_id = project_id or self.project_id
+        parent = f"projects/{project_id}/locations/{location}"
+        dataset = f"projects/{project_id}/locations/{location}/datasets/{dataset_id}"
+        result = client.create_model(
+            request={
+                "parent": parent,
+                "model": {
+                    "display_name": display_name,
+                    "dataset": dataset,
+                },
+            },
+            retry=retry,
+            timeout=timeout,
+            metadata=metadata,
+        )
+        return result
+
+    def get_model(
+        self,
+        model_id: str,
+        project_id: str,
+        location: str,
+        retry: Retry | _MethodDefault = DEFAULT,
+        timeout: float | _MethodDefault = DEFAULT,
+        metadata: Sequence[tuple[str, str]] = (),
+    ) -> automl_translation.Model:
+        """
+        Retrieve the dataset for the given model_id.
+
+        :param model_id: ID of translation model to be retrieved.
+        :param project_id: ID of the Google Cloud project where dataset is located. If not provided
+            default project_id is used.
+        :param location: The location of the project.
+        :param retry: A retry object used to retry requests. If `None` is specified, requests will not be
+            retried.
+        :param timeout: The amount of time, in seconds, to wait for the request to complete. Note that if
+            `retry` is specified, the timeout applies to each individual attempt.
+        :param metadata: Additional metadata that is provided to the method.
+
+        :return: `automl_translation.Model` instance.
+        """
+        client = self.get_client()
+        name = f"projects/{project_id}/locations/{location}/models/{model_id}"
+        return client.get_model(
+            request={"name": name},
+            retry=retry,
+            timeout=timeout,
+            metadata=metadata,
+        )
+
+    def list_models(
+        self,
+        project_id: str,
+        location: str,
+        filter_str: str | None = None,
+        page_size: int | None = None,
+        retry: Retry | _MethodDefault = DEFAULT,
+        timeout: float | _MethodDefault = DEFAULT,
+        metadata: Sequence[tuple[str, str]] = (),
+    ) -> pagers.ListModelsPager:
+        """
+        List translation models in a project.
+
+        :param project_id: ID of the Google Cloud project where models are located. If not provided
+            default project_id is used.
+        :param location: The location of the project.
+        :param filter_str: An optional expression for filtering the models that will
+            be returned. Supported filter: ``dataset_id=${dataset_id}``.
+        :param page_size: Optional custom page size value. The server can
+            return fewer results than requested.
+        :param retry: A retry object used to retry requests. If `None` is specified, requests will not be
+            retried.
+        :param timeout: The amount of time, in seconds, to wait for the request to complete. Note that if
+            `retry` is specified, the timeout applies to each individual attempt.
+        :param metadata: Additional metadata that is provided to the method.
+
+        :return: ``pagers.ListDatasetsPager`` instance, iterable object to retrieve the datasets list.
+        """
+        client = self.get_client()
+        parent = f"projects/{project_id}/locations/{location}"
+        result = client.list_models(
+            request={
+                "parent": parent,
+                "filter": filter_str,
+                "page_size": page_size,
+            },
+            retry=retry,
+            timeout=timeout,
+            metadata=metadata,
+        )
+        return result
+
+    def delete_model(
+        self,
+        model_id: str,
+        project_id: str,
+        location: str,
+        retry: Retry | _MethodDefault = DEFAULT,
+        timeout: float | None = None,
+        metadata: Sequence[tuple[str, str]] = (),
+    ) -> Operation:
+        """
+        Delete the translation model and all of its contents.
+
+        :param model_id: ID of model to be deleted.
+        :param project_id: ID of the Google Cloud project where dataset is located. If not provided
+            default project_id is used.
+        :param location: The location of the project.
+        :param retry: A retry object used to retry requests. If `None` is specified, requests will not be
+            retried.
+        :param timeout: The amount of time, in seconds, to wait for the request to complete. Note that if
+            `retry` is specified, the timeout applies to each individual attempt.
+        :param metadata: Additional metadata that is provided to the method.
+
+        :return: `Operation` object with dataset deletion results, when finished.
+        """
+        client = self.get_client()
+        name = f"projects/{project_id}/locations/{location}/models/{model_id}"
+        result = client.delete_model(
+            request={"name": name},
+            retry=retry,
+            timeout=timeout,
+            metadata=metadata,
+        )
+        return result
+
+    def translate_document(
+        self,
+        *,
+        project_id: str = PROVIDE_PROJECT_ID,
+        source_language_code: str | None = None,
+        target_language_code: str,
+        location: str | None = None,
+        document_input_config: DocumentInputConfig | dict,
+        document_output_config: DocumentOutputConfig | dict | None,
+        customized_attribution: str | None = None,
+        is_translate_native_pdf_only: bool = False,
+        enable_shadow_removal_native_pdf: bool = False,
+        enable_rotation_correction: bool = False,
+        model: str | None = None,
+        glossary_config: TranslateTextGlossaryConfig | None = None,
+        labels: str | None = None,
+        timeout: float | _MethodDefault = DEFAULT,
+        metadata: Sequence[tuple[str, str]] = (),
+        retry: Retry | _MethodDefault | None = DEFAULT,
+    ) -> TranslateDocumentResponse:
+        """
+        Translate the document provided.
+
+        :param project_id: Required. The ID of the Google Cloud project that the service belongs to.
+        :param source_language_code: Optional. The ISO-639 language code of the
+            input document text if known. If the source language isn't specified,
+            the API attempts to identify the source language automatically and returns
+            the source language within the response.
+        :param target_language_code: Required. The ISO-639 language code to use
+            for translation of the input document text.
+        :param location: Optional. Project or location to make a call. Must refer to
+            a caller's project.
+            If not specified, 'global' is used.
+            Non-global location is required for requests using AutoML models or custom glossaries.
+            Models and glossaries must be within the same region (have the same location-id).
+        :param document_input_config: A document translation request input config.
+        :param document_output_config: Optional. A document translation request output config.
+            If not provided the translated file will only be returned through a byte-stream
+            and its output mime type will be the same as the input file's mime type.
+        :param customized_attribution: Optional. This flag is to support user customized
+            attribution. If not provided, the default is ``Machine Translated by Google``.
+            Customized attribution should follow rules in
+            https://cloud.google.com/translate/attribution#attribution_and_logos
+        :param is_translate_native_pdf_only: Optional. Param for external
+            customers. If true, the page limit of online native PDF
+            translation is 300 and only native PDF pages will be
+            translated.
+        :param enable_shadow_removal_native_pdf: Optional. If true, use the text removal server to remove the
+            shadow text on background image for native PDF translation.
+            Shadow removal feature can only be enabled when both ``is_translate_native_pdf_only``,
+            ``pdf_native_only`` are False.
+        :param enable_rotation_correction: Optional. If true, enable auto rotation
+            correction in DVS.
+        :param model: Optional. The ``model`` type requested for this translation.
+            If not provided, the default Google model (NMT) will be used.
+            The format depends on model type:
+
+            -  AutoML Translation models:
+               ``projects/{project-number-or-id}/locations/{location-id}/models/{model-id}``
+            -  General (built-in) models:
+               ``projects/{project-number-or-id}/locations/{location-id}/models/general/nmt``,
+
+            If not provided, the default Google model (NMT) will be used
+            for translation.
+        :param glossary_config: Optional. Glossary to be applied. The glossary must be
+            within the same region (have the same location-id) as the
+            model.
+        :param labels: Optional. The labels with user-defined
+            metadata for the request.
+            See https://cloud.google.com/translate/docs/advanced/labels for more information.
+        :param retry: Designation of what errors, if any, should be retried.
+        :param timeout: The timeout for this request.
+        :param metadata: Strings which should be sent along with the request as metadata.
+
+        :return: Translate document result from the API response.
+        """
+        client = self.get_client()
+        location_id = "global" if not location else location
+        parent = f"projects/{project_id or self.project_id}/locations/{location_id}"
+        return client.translate_document(
+            request={
+                "parent": parent,
+                "source_language_code": source_language_code,
+                "target_language_code": target_language_code,
+                "document_input_config": document_input_config,
+                "document_output_config": document_output_config,
+                "customized_attribution": customized_attribution,
+                "is_translate_native_pdf_only": is_translate_native_pdf_only,
+                "enable_shadow_removal_native_pdf": enable_shadow_removal_native_pdf,
+                "enable_rotation_correction": enable_rotation_correction,
+                "model": model,
+                "glossary_config": glossary_config,
+                "labels": labels,
+            },
+            timeout=timeout,
+            retry=retry,
+            metadata=metadata,
+        )
+
+    def batch_translate_document(
+        self,
+        *,
+        project_id: str = PROVIDE_PROJECT_ID,
+        source_language_code: str,
+        target_language_codes: MutableSequence[str] | None = None,
+        location: str | None = None,
+        input_configs: MutableSequence[BatchDocumentInputConfig | dict],
+        output_config: BatchDocumentOutputConfig | dict,
+        customized_attribution: str | None = None,
+        format_conversions: MutableMapping[str, str] | None = None,
+        enable_shadow_removal_native_pdf: bool = False,
+        enable_rotation_correction: bool = False,
+        models: MutableMapping[str, str] | None = None,
+        glossaries: MutableMapping[str, TranslateTextGlossaryConfig] | None = None,
+        timeout: float | _MethodDefault = DEFAULT,
+        metadata: Sequence[tuple[str, str]] = (),
+        retry: Retry | _MethodDefault | None = DEFAULT,
+    ) -> Operation:
+        """
+        Translate documents batch by configs provided.
+
+        :param project_id: Required. The ID of the Google Cloud project that the service belongs to.
+        :param source_language_code: Optional. The ISO-639 language code of the
+            input text if known. If the source language isn't specified, the API attempts to identify
+            the source language automatically and returns the source language within the response.
+        :param target_language_codes: Required. The ISO-639 language code to use
+            for translation of the input document. Specify up to 10 language codes here.
+        :param location: Optional. Project or location to make a call. Must refer to
+            a caller's project. If not specified, 'global' is used.
+            Non-global location is required for requests using AutoML models or custom glossaries.
+            Models and glossaries must be within the same region (have the same location-id).
+        :param input_configs: Input configurations. The total number of files matched should be <=
+            100. The total content size to translate should be <= 100M Unicode codepoints.
+            The files must use UTF-8 encoding.
+        :param output_config: Output configuration. If 2 input configs match to the same file (that
+            is, same input path), no output for duplicate inputs will be generated.
+        :param format_conversions: Optional. The file format conversion map that is applied to
+            all input files. The map key is the original mime_type.
+            The map value is the target mime_type of translated documents.
+            Supported file format conversion includes:
+
+            -  ``application/pdf`` to
+               ``application/vnd.openxmlformats-officedocument.wordprocessingml.document``
+
+            If nothing specified, output files will be in the same format as the original file.
+        :param customized_attribution: Optional. This flag is to support user customized
+            attribution. If not provided, the default is ``Machine Translated by Google``.
+            Customized attribution should follow rules in
+            https://cloud.google.com/translate/attribution#attribution_and_logos
+        :param enable_shadow_removal_native_pdf: Optional. If true, use the text removal server to remove the
+            shadow text on background image for native PDF translation.
+            Shadow removal feature can only be enabled when both ``is_translate_native_pdf_only``,
+            ``pdf_native_only`` are False.
+        :param enable_rotation_correction: Optional. If true, enable auto rotation
+            correction in DVS.
+        :param models: Optional. The models to use for translation. Map's key is
+            target language code. Map's value is the model name. Value
+            can be a built-in general model, or an AutoML Translation model.
+            The value format depends on model type:
+
+            -  AutoML Translation models:
+               ``projects/{project-number-or-id}/locations/{location-id}/models/{model-id}``
+            -  General (built-in) models:
+               ``projects/{project-number-or-id}/locations/{location-id}/models/general/nmt``,
+
+            If the map is empty or a specific model is not requested for
+            a language pair, then default google model (NMT) is used.
+        :param glossaries: Glossaries to be applied. It's keyed by target language code.
+        :param retry: Designation of what errors, if any, should be retried.
+        :param timeout: The timeout for this request.
+        :param metadata: Strings which should be sent along with the request as metadata.
+
+        :return: Batch translate document result from the API response.
+        """
+        client = self.get_client()
+        location_id = "global" if not location else location
+        parent = f"projects/{project_id or self.project_id}/locations/{location_id}"
+        return client.batch_translate_document(
+            request={
+                "parent": parent,
+                "source_language_code": source_language_code,
+                "target_language_codes": target_language_codes,
+                "input_configs": input_configs,
+                "output_config": output_config,
+                "format_conversions": format_conversions,
+                "customized_attribution": customized_attribution,
+                "enable_shadow_removal_native_pdf": enable_shadow_removal_native_pdf,
+                "enable_rotation_correction": enable_rotation_correction,
+                "models": models,
+                "glossaries": glossaries,
+            },
+            timeout=timeout,
+            retry=retry,
+            metadata=metadata,
+        )

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

@@ -0,0 +1,147 @@
+# 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.
+"""This module contains a Google Cloud Vertex AI Feature Store hook."""
+
+from __future__ import annotations
+
+from google.api_core.client_options import ClientOptions
+from google.cloud.aiplatform_v1beta1 import (
+    FeatureOnlineStoreAdminServiceClient,
+)
+
+from airflow.exceptions import AirflowException
+from airflow.providers.google.common.consts import CLIENT_INFO
+from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID, GoogleBaseHook
+
+
+class FeatureStoreHook(GoogleBaseHook):
+    """
+    Hook for interacting with Google Cloud Vertex AI Feature Store.
+
+    This hook provides an interface to manage Feature Store resources in Vertex AI,
+    including feature views and their synchronization operations. It handles authentication
+    and provides methods for common Feature Store operations.
+
+    :param gcp_conn_id: The connection ID to use for connecting to Google Cloud Platform.
+        Defaults to 'google_cloud_default'.
+    :param impersonation_chain: Optional service account to impersonate using short-term
+        credentials. Can be either a single account or a chain of accounts required to
+        get the access_token of the last account in the list, which will be impersonated
+        in the request. If set as a string, the account must grant the originating account
+        the Service Account Token Creator IAM role. 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.
+    """
+
+    def get_feature_online_store_admin_service_client(
+        self,
+        location: str | None = None,
+    ) -> FeatureOnlineStoreAdminServiceClient:
+        """
+        Create and returns a FeatureOnlineStoreAdminServiceClient object.
+
+        This method initializes a client for interacting with the Feature Store API,
+        handling proper endpoint configuration based on the specified location.
+
+        :param location: Optional. The Google Cloud region where the service is located.
+            If provided and not 'global', the client will be configured to use the
+            region-specific API endpoint.
+        """
+        if location and location != "global":
+            client_options = ClientOptions(api_endpoint=f"{location}-aiplatform.googleapis.com:443")
+        else:
+            client_options = ClientOptions()
+        return FeatureOnlineStoreAdminServiceClient(
+            credentials=self.get_credentials(), client_info=CLIENT_INFO, client_options=client_options
+        )
+
+    def get_feature_view_sync(
+        self,
+        location: str,
+        feature_view_sync_name: str,
+    ) -> dict:
+        """
+        Retrieve the status and details of a Feature View synchronization operation.
+
+        This method fetches information about a specific feature view sync operation,
+        including its current status, timing information, and synchronization metrics.
+
+        :param location: The Google Cloud region where the feature store is located
+            (e.g., 'us-central1', 'us-east1').
+        :param feature_view_sync_name: The full resource name of the feature view
+            sync operation to retrieve.
+        """
+        client = self.get_feature_online_store_admin_service_client(location)
+
+        try:
+            response = client.get_feature_view_sync(name=feature_view_sync_name)
+
+            report = {
+                "name": feature_view_sync_name,
+                "start_time": int(response.run_time.start_time.seconds),
+            }
+
+            if hasattr(response.run_time, "end_time") and response.run_time.end_time.seconds:
+                report["end_time"] = int(response.run_time.end_time.seconds)
+                report["sync_summary"] = {
+                    "row_synced": int(response.sync_summary.row_synced),
+                    "total_slot": int(response.sync_summary.total_slot),
+                }
+
+            return report
+
+        except Exception as e:
+            self.log.error("Failed to get feature view sync: %s", str(e))
+            raise AirflowException(str(e))
+
+    @GoogleBaseHook.fallback_to_default_project_id
+    def sync_feature_view(
+        self,
+        location: str,
+        feature_online_store_id: str,
+        feature_view_id: str,
+        project_id: str = PROVIDE_PROJECT_ID,
+    ) -> str:
+        """
+        Initiate a synchronization operation for a Feature View.
+
+        This method triggers a sync operation that updates the online serving data
+        for a feature view based on the latest data in the underlying batch source.
+        The sync operation ensures that the online feature values are up-to-date
+        for real-time serving.
+
+        :param location: The Google Cloud region where the feature store is located
+            (e.g., 'us-central1', 'us-east1').
+        :param feature_online_store_id: The ID of the online feature store that
+            contains the feature view to be synchronized.
+        :param feature_view_id: The ID of the feature view to synchronize.
+        :param project_id: The ID of the Google Cloud project that contains the
+            feature store. If not provided, will attempt to determine from the
+            environment.
+        """
+        client = self.get_feature_online_store_admin_service_client(location)
+        feature_view = f"projects/{project_id}/locations/{location}/featureOnlineStores/{feature_online_store_id}/featureViews/{feature_view_id}"
+
+        try:
+            response = client.sync_feature_view(feature_view=feature_view)
+
+            return str(response.feature_view_sync)
+
+        except Exception as e:
+            self.log.error("Failed to sync feature view: %s", str(e))
+            raise AirflowException(str(e))

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

@@ -43,6 +43,11 @@ if TYPE_CHECKING:
 class GenerativeModelHook(GoogleBaseHook):
     """Hook for Google Cloud Vertex AI Generative Model APIs."""
 
+    @deprecated(
+        planned_removal_date="April 09, 2025",
+        use_instead="GenerativeModelHook.get_generative_model",
+        category=AirflowProviderDeprecationWarning,
+    )
     def get_text_generation_model(self, pretrained_model: str):
         """Return a Model Garden Model object based on Text Generation."""
         model = TextGenerationModel.from_pretrained(pretrained_model)
@@ -275,6 +280,11 @@ class GenerativeModelHook(GoogleBaseHook):
 
         return response.text
 
+    @deprecated(
+        planned_removal_date="April 09, 2025",
+        use_instead="GenerativeModelHook.generative_model_generate_content",
+        category=AirflowProviderDeprecationWarning,
+    )
     @GoogleBaseHook.fallback_to_default_project_id
     def text_generation_model_predict(
         self,

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

@@ -0,0 +1,55 @@
+#
+# 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.
+"""This module contains Google Cloud AlloyDB links."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from airflow.providers.google.cloud.links.base import BaseGoogleLink
+
+if TYPE_CHECKING:
+    from airflow.models import BaseOperator
+    from airflow.utils.context import Context
+
+ALLOY_DB_BASE_LINK = "/alloydb"
+ALLOY_DB_CLUSTER_LINK = (
+    ALLOY_DB_BASE_LINK + "/locations/{location_id}/clusters/{cluster_id}?project={project_id}"
+)
+
+
+class AlloyDBClusterLink(BaseGoogleLink):
+    """Helper class for constructing AlloyDB cluster Link."""
+
+    name = "AlloyDB Cluster"
+    key = "alloy_db_cluster"
+    format_str = ALLOY_DB_CLUSTER_LINK
+
+    @staticmethod
+    def persist(
+        context: Context,
+        task_instance: BaseOperator,
+        location_id: str,
+        cluster_id: str,
+        project_id: str | None,
+    ):
+        task_instance.xcom_push(
+            context,
+            key=AlloyDBClusterLink.key,
+            value={"location_id": location_id, "cluster_id": cluster_id, "project_id": project_id},
+        )