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

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

@@ -20,7 +20,7 @@
 from __future__ import annotations
 
 from collections.abc import MutableMapping, MutableSequence, Sequence
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, cast
 
 from google.api_core.exceptions import GoogleAPICallError
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
@@ -28,8 +28,11 @@ from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
 from airflow.exceptions import AirflowException
 from airflow.providers.google.cloud.hooks.translate import CloudTranslateHook, TranslateHook
 from airflow.providers.google.cloud.links.translate import (
+    TranslateResultByOutputConfigLink,
     TranslateTextBatchLink,
     TranslationDatasetsListLink,
+    TranslationModelLink,
+    TranslationModelsListLink,
     TranslationNativeDatasetLink,
 )
 from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
@@ -38,7 +41,11 @@ from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
 if TYPE_CHECKING:
     from google.api_core.retry import Retry
     from google.cloud.translate_v3.types import (
+        BatchDocumentInputConfig,
+        BatchDocumentOutputConfig,
         DatasetInputConfig,
+        DocumentInputConfig,
+        DocumentOutputConfig,
         InputConfig,
         OutputConfig,
         TranslateTextGlossaryConfig,
@@ -723,3 +730,581 @@ class TranslateDeleteDatasetOperator(GoogleCloudBaseOperator):
         )
         hook.wait_for_operation_done(operation=operation, timeout=self.timeout)
         self.log.info("Dataset deletion complete!")
+
+
+class TranslateCreateModelOperator(GoogleCloudBaseOperator):
+    """
+    Creates a Google Cloud Translate model.
+
+    Creates a `native` translation model, using API V3.
+    For more information on how to use this operator, take a look at the guide:
+    :ref:`howto/operator:TranslateCreateModelOperator`.
+
+    :param dataset_id: The dataset id used for model training.
+    :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: 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.
+    :param gcp_conn_id: The connection ID to use connecting to Google Cloud.
+    :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.
+        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 (templated).
+    """
+
+    template_fields: Sequence[str] = (
+        "dataset_id",
+        "location",
+        "project_id",
+        "gcp_conn_id",
+        "impersonation_chain",
+    )
+
+    operator_extra_links = (TranslationModelLink(),)
+
+    def __init__(
+        self,
+        *,
+        project_id: str = PROVIDE_PROJECT_ID,
+        location: str,
+        dataset_id: str,
+        display_name: str,
+        timeout: float | None = None,
+        retry: Retry | _MethodDefault = DEFAULT,
+        gcp_conn_id: str = "google_cloud_default",
+        metadata: Sequence[tuple[str, str]] = (),
+        impersonation_chain: str | Sequence[str] | None = None,
+        **kwargs,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.project_id = project_id
+        self.location = location
+        self.dataset_id = dataset_id
+        self.display_name = display_name
+        self.metadata = metadata
+        self.timeout = timeout
+        self.retry = retry
+        self.gcp_conn_id = gcp_conn_id
+        self.impersonation_chain = impersonation_chain
+
+    def execute(self, context: Context) -> str:
+        hook = TranslateHook(
+            gcp_conn_id=self.gcp_conn_id,
+            impersonation_chain=self.impersonation_chain,
+        )
+        self.log.info("Model creation started, dataset_id %s...", self.dataset_id)
+        try:
+            result_operation = hook.create_model(
+                dataset_id=self.dataset_id,
+                display_name=self.display_name,
+                location=self.location,
+                project_id=self.project_id,
+                retry=self.retry,
+                timeout=self.timeout,
+                metadata=self.metadata,
+            )
+        except GoogleAPICallError as e:
+            self.log.error("Error submitting create_model operation ")
+            raise AirflowException(e)
+
+        self.log.info("Training has started")
+        hook.wait_for_operation_done(operation=result_operation)
+        result = hook.wait_for_operation_result(operation=result_operation)
+        result = type(result).to_dict(result)
+        model_id = hook.extract_object_id(result)
+        self.xcom_push(context, key="model_id", value=model_id)
+        self.log.info("Model creation complete. The model_id: %s.", model_id)
+
+        project_id = self.project_id or hook.project_id
+        TranslationModelLink.persist(
+            context=context,
+            task_instance=self,
+            dataset_id=self.dataset_id,
+            model_id=model_id,
+            project_id=project_id,
+        )
+        return result
+
+
+class TranslateModelsListOperator(GoogleCloudBaseOperator):
+    """
+    Get a list of native Google Cloud Translation models in a project.
+
+    Get project's list of `native` translation models, using API V3.
+    For more information on how to use this operator, take a look at the guide:
+    :ref:`howto/operator:TranslateModelsListOperator`.
+
+    :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: 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.
+    :param gcp_conn_id: The connection ID to use connecting to Google Cloud.
+    :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.
+        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 (templated).
+    """
+
+    template_fields: Sequence[str] = (
+        "location",
+        "project_id",
+        "gcp_conn_id",
+        "impersonation_chain",
+    )
+
+    operator_extra_links = (TranslationModelsListLink(),)
+
+    def __init__(
+        self,
+        *,
+        project_id: str = PROVIDE_PROJECT_ID,
+        location: str,
+        metadata: Sequence[tuple[str, str]] = (),
+        timeout: float | _MethodDefault = DEFAULT,
+        retry: Retry | _MethodDefault = DEFAULT,
+        gcp_conn_id: str = "google_cloud_default",
+        impersonation_chain: str | Sequence[str] | None = None,
+        **kwargs,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.project_id = project_id
+        self.location = location
+        self.metadata = metadata
+        self.timeout = timeout
+        self.retry = retry
+        self.gcp_conn_id = gcp_conn_id
+        self.impersonation_chain = impersonation_chain
+
+    def execute(self, context: Context):
+        hook = TranslateHook(
+            gcp_conn_id=self.gcp_conn_id,
+            impersonation_chain=self.impersonation_chain,
+        )
+        project_id = self.project_id or hook.project_id
+        TranslationModelsListLink.persist(
+            context=context,
+            task_instance=self,
+            project_id=project_id,
+        )
+        self.log.info("Requesting models list")
+        results_pager = hook.list_models(
+            location=self.location,
+            project_id=self.project_id,
+            retry=self.retry,
+            timeout=self.timeout,
+            metadata=self.metadata,
+        )
+        result_ids = []
+        for model_item in results_pager:
+            model_data = type(model_item).to_dict(model_item)
+            model_id = hook.extract_object_id(model_data)
+            result_ids.append(model_id)
+        self.log.info("Fetching the models list complete. Model id-s: %s", result_ids)
+        return result_ids
+
+
+class TranslateDeleteModelOperator(GoogleCloudBaseOperator):
+    """
+    Delete translation model and all of its contents.
+
+    Deletes the translation model and it's data, using API V3.
+    For more information on how to use this operator, take a look at the guide:
+    :ref:`howto/operator:TranslateDeleteModelOperator`.
+
+    :param model_id: The model_id of target native model to be deleted.
+    :param location: The location of the project.
+    :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.
+    :param gcp_conn_id: The connection ID to use connecting to Google Cloud.
+    :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.
+        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 (templated).
+    """
+
+    template_fields: Sequence[str] = (
+        "model_id",
+        "location",
+        "project_id",
+        "gcp_conn_id",
+        "impersonation_chain",
+    )
+
+    def __init__(
+        self,
+        *,
+        model_id: str,
+        location: str,
+        project_id: str = PROVIDE_PROJECT_ID,
+        metadata: Sequence[tuple[str, str]] = (),
+        timeout: float | None = None,
+        retry: Retry | _MethodDefault = DEFAULT,
+        gcp_conn_id: str = "google_cloud_default",
+        impersonation_chain: str | Sequence[str] | None = None,
+        **kwargs,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.model_id = model_id
+        self.project_id = project_id
+        self.location = location
+        self.metadata = metadata
+        self.timeout = timeout
+        self.retry = retry
+        self.gcp_conn_id = gcp_conn_id
+        self.impersonation_chain = impersonation_chain
+
+    def execute(self, context: Context):
+        hook = TranslateHook(gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain)
+        self.log.info("Deleting the model %s...", self.model_id)
+        operation = hook.delete_model(
+            model_id=self.model_id,
+            location=self.location,
+            project_id=self.project_id,
+            retry=self.retry,
+            timeout=self.timeout,
+            metadata=self.metadata,
+        )
+        hook.wait_for_operation_done(operation=operation, timeout=self.timeout)
+        self.log.info("Model deletion complete!")
+
+
+class TranslateDocumentOperator(GoogleCloudBaseOperator):
+    """
+    Translate document provided.
+
+    Wraps the Google cloud Translate Text (Advanced) functionality.
+    Supports wide range of input/output file types, please visit the
+    https://cloud.google.com/translate/docs/advanced/translate-documents for more details.
+
+    For more information on how to use this operator, take a look at the guide:
+        :ref:`howto/operator:TranslateDocumentOperator`.
+
+    :param project_id: Optional. The ID of the Google Cloud project that the
+        service belongs to. If not specified the hook project_id will be used.
+    :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.
+    :param transliteration_config: Optional. Transliteration to be applied.
+    :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.
+    :param gcp_conn_id: The connection ID to use connecting to Google Cloud.
+    :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.
+        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 (templated).
+    """
+
+    operator_extra_links = (TranslateResultByOutputConfigLink(),)
+
+    template_fields: Sequence[str] = (
+        "source_language_code",
+        "target_language_code",
+        "document_input_config",
+        "document_output_config",
+        "model",
+        "gcp_conn_id",
+        "impersonation_chain",
+    )
+
+    def __init__(
+        self,
+        *,
+        location: str | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
+        source_language_code: str | None = None,
+        target_language_code: str,
+        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,
+        retry: Retry | _MethodDefault | None = DEFAULT,
+        metadata: Sequence[tuple[str, str]] = (),
+        gcp_conn_id: str = "google_cloud_default",
+        impersonation_chain: str | Sequence[str] | None = None,
+        **kwargs,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.project_id = project_id
+        self.source_language_code = source_language_code
+        self.target_language_code = target_language_code
+        self.document_input_config = document_input_config
+        self.document_output_config = document_output_config
+        self.customized_attribution = customized_attribution
+        self.is_translate_native_pdf_only = is_translate_native_pdf_only
+        self.enable_shadow_removal_native_pdf = enable_shadow_removal_native_pdf
+        self.enable_rotation_correction = enable_rotation_correction
+        self.location = location
+        self.labels = labels
+        self.model = model
+        self.glossary_config = glossary_config
+        self.metadate = metadata
+        self.timeout = timeout
+        self.retry = retry
+        self.gcp_conn_id = gcp_conn_id
+        self.impersonation_chain = impersonation_chain
+
+    def execute(self, context: Context) -> dict:
+        hook = TranslateHook(
+            gcp_conn_id=self.gcp_conn_id,
+            impersonation_chain=self.impersonation_chain,
+        )
+        try:
+            self.log.info("Starting the document translation")
+            doc_translation_result = hook.translate_document(
+                source_language_code=self.source_language_code,
+                target_language_code=self.target_language_code,
+                document_input_config=self.document_input_config,
+                document_output_config=self.document_output_config,
+                customized_attribution=self.customized_attribution,
+                is_translate_native_pdf_only=self.is_translate_native_pdf_only,
+                enable_shadow_removal_native_pdf=self.enable_shadow_removal_native_pdf,
+                enable_rotation_correction=self.enable_rotation_correction,
+                location=self.location,
+                labels=self.labels,
+                model=self.model,
+                glossary_config=self.glossary_config,
+                timeout=self.timeout,
+                retry=self.retry,
+                metadata=self.metadate,
+            )
+            self.log.info("Document translation completed")
+        except GoogleAPICallError as e:
+            self.log.error("An error occurred executing translate_document method: \n%s", e)
+            raise AirflowException(e)
+        if self.document_output_config:
+            TranslateResultByOutputConfigLink.persist(
+                context=context,
+                task_instance=self,
+                project_id=self.project_id or hook.project_id,
+                output_config=self.document_output_config,
+            )
+        return cast(dict, type(doc_translation_result).to_dict(doc_translation_result))
+
+
+class TranslateDocumentBatchOperator(GoogleCloudBaseOperator):
+    """
+    Translate documents provided via input and output configurations.
+
+    Up to 10 target languages per operation supported.
+    Wraps the Google cloud Translate Text (Advanced) functionality.
+    See https://cloud.google.com/translate/docs/advanced/batch-translation.
+
+    For more information on how to use this operator, take a look at the guide:
+    :ref:`howto/operator:TranslateDocumentBatchOperator`.
+
+    :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.
+    :param gcp_conn_id: The connection ID to use connecting to Google Cloud.
+    :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.
+        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 (templated).
+    """
+
+    operator_extra_links = (TranslateResultByOutputConfigLink(),)
+
+    template_fields: Sequence[str] = (
+        "input_configs",
+        "output_config",
+        "target_language_codes",
+        "source_language_code",
+        "models",
+        "glossaries",
+        "gcp_conn_id",
+        "impersonation_chain",
+    )
+
+    def __init__(
+        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,
+        metadata: Sequence[tuple[str, str]] = (),
+        timeout: float | _MethodDefault = DEFAULT,
+        retry: Retry | _MethodDefault | None = DEFAULT,
+        gcp_conn_id: str = "google_cloud_default",
+        impersonation_chain: str | Sequence[str] | None = None,
+        **kwargs,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.project_id = project_id
+        self.location = location
+        self.target_language_codes = target_language_codes
+        self.source_language_code = source_language_code
+        self.input_configs = input_configs
+        self.output_config = output_config
+        self.customized_attribution = customized_attribution
+        self.format_conversions = format_conversions
+        self.enable_shadow_removal_native_pdf = enable_shadow_removal_native_pdf
+        self.enable_rotation_correction = enable_rotation_correction
+        self.models = models
+        self.glossaries = glossaries
+        self.metadata = metadata
+        self.timeout = timeout
+        self.retry = retry
+        self.gcp_conn_id = gcp_conn_id
+        self.impersonation_chain = impersonation_chain
+
+    def execute(self, context: Context) -> dict:
+        hook = TranslateHook(
+            gcp_conn_id=self.gcp_conn_id,
+            impersonation_chain=self.impersonation_chain,
+        )
+        try:
+            batch_document_translate_operation = hook.batch_translate_document(
+                project_id=self.project_id,
+                location=self.location,
+                target_language_codes=self.target_language_codes,
+                source_language_code=self.source_language_code,
+                input_configs=self.input_configs,
+                output_config=self.output_config,
+                customized_attribution=self.customized_attribution,
+                format_conversions=self.format_conversions,
+                enable_shadow_removal_native_pdf=self.enable_shadow_removal_native_pdf,
+                enable_rotation_correction=self.enable_rotation_correction,
+                models=self.models,
+                glossaries=self.glossaries,
+                metadata=self.metadata,
+                timeout=self.timeout,
+                retry=self.retry,
+            )
+        except GoogleAPICallError as e:
+            self.log.error("An error occurred executing batch_translate_document method: \n%s", e)
+            raise AirflowException(e)
+        self.log.info("Batch document translation job started.")
+        TranslateResultByOutputConfigLink.persist(
+            context=context,
+            task_instance=self,
+            project_id=self.project_id or hook.project_id,
+            output_config=self.output_config,
+        )
+        result = hook.wait_for_operation_result(batch_document_translate_operation)
+        self.log.info("Batch document translation job finished")
+        return cast(dict, type(result).to_dict(result))