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

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

@@ -33,12 +33,6 @@ from functools import cached_property
 from typing import TYPE_CHECKING, Any
 from uuid import uuid4
 
-from google.api_core.exceptions import AlreadyExists, GoogleAPICallError
-from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
-from google.cloud.exceptions import NotFound
-from google.cloud.pubsub_v1 import PublisherClient, SubscriberClient
-from google.cloud.pubsub_v1.types import PublisherOptions
-from google.pubsub_v1.services.subscriber.async_client import SubscriberAsyncClient
 from googleapiclient.errors import HttpError
 
 from airflow.exceptions import AirflowProviderDeprecationWarning
@@ -49,6 +43,12 @@ from airflow.providers.google.common.hooks.base_google import (
     GoogleBaseHook,
 )
 from airflow.version import version
+from google.api_core.exceptions import AlreadyExists, GoogleAPICallError
+from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
+from google.cloud.exceptions import NotFound
+from google.cloud.pubsub_v1 import PublisherClient, SubscriberClient
+from google.cloud.pubsub_v1.types import PublisherOptions
+from google.pubsub_v1.services.subscriber.async_client import SubscriberAsyncClient
 
 if TYPE_CHECKING:
     from google.api_core.retry import Retry

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

@@ -23,6 +23,8 @@ from collections.abc import Sequence
 from functools import cached_property
 from typing import TYPE_CHECKING
 
+from airflow.providers.google.common.consts import CLIENT_INFO
+from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
 from google.cloud.secretmanager_v1 import (
     AccessSecretVersionResponse,
@@ -32,9 +34,6 @@ from google.cloud.secretmanager_v1 import (
     SecretVersion,
 )
 
-from airflow.providers.google.common.consts import CLIENT_INFO
-from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
-
 if TYPE_CHECKING:
     from google.api_core.retry import Retry
     from google.cloud.secretmanager_v1.services.secret_manager_service.pagers import ListSecretsPager

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

@@ -22,14 +22,14 @@ from __future__ import annotations
 from collections.abc import Sequence
 from typing import TYPE_CHECKING, Callable, NamedTuple
 
-from google.api_core.exceptions import AlreadyExists, GoogleAPICallError
-from google.cloud.spanner_v1.client import Client
 from sqlalchemy import create_engine
 
 from airflow.exceptions import AirflowException
 from airflow.providers.common.sql.hooks.sql import DbApiHook
 from airflow.providers.google.common.consts import CLIENT_INFO
 from airflow.providers.google.common.hooks.base_google import GoogleBaseHook, get_field
+from google.api_core.exceptions import AlreadyExists, GoogleAPICallError
+from google.cloud.spanner_v1.client import Client
 
 if TYPE_CHECKING:
     from google.cloud.spanner_v1.database import Database

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

@@ -22,13 +22,12 @@ from __future__ import annotations
 from collections.abc import Sequence
 from typing import TYPE_CHECKING
 
+from airflow.providers.google.common.consts import CLIENT_INFO
+from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
 from google.cloud.speech_v1 import SpeechClient
 from google.cloud.speech_v1.types import RecognitionAudio, RecognitionConfig
 
-from airflow.providers.google.common.consts import CLIENT_INFO
-from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
-
 if TYPE_CHECKING:
     from google.api_core.retry import Retry
 

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

@@ -24,15 +24,15 @@ import json
 from collections.abc import Sequence
 from typing import TYPE_CHECKING, Any
 
+from googleapiclient.errors import HttpError
+
+from airflow.exceptions import AirflowException
+from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID, GoogleBaseHook
 from google.api_core.exceptions import InvalidArgument
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
 from google.cloud import monitoring_v3
 from google.cloud.monitoring_v3 import AlertPolicy, NotificationChannel
 from google.protobuf.field_mask_pb2 import FieldMask
-from googleapiclient.errors import HttpError
-
-from airflow.exceptions import AirflowException
-from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID, GoogleBaseHook
 
 if TYPE_CHECKING:
     from google.api_core.retry import Retry

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

@@ -22,13 +22,12 @@ from __future__ import annotations
 from collections.abc import Sequence
 from typing import TYPE_CHECKING
 
-from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
-from google.cloud.tasks_v2 import CloudTasksClient
-from google.cloud.tasks_v2.types import Queue, Task
-
 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
+from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
+from google.cloud.tasks_v2 import CloudTasksClient
+from google.cloud.tasks_v2.types import Queue, Task
 
 if TYPE_CHECKING:
     from google.api_core.retry import Retry

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

@@ -22,6 +22,8 @@ from __future__ import annotations
 from collections.abc import Sequence
 from typing import TYPE_CHECKING
 
+from airflow.providers.google.common.consts import CLIENT_INFO
+from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
 from google.cloud.texttospeech_v1 import TextToSpeechClient
 from google.cloud.texttospeech_v1.types import (
@@ -31,9 +33,6 @@ from google.cloud.texttospeech_v1.types import (
     VoiceSelectionParams,
 )
 
-from airflow.providers.google.common.consts import CLIENT_INFO
-from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
-
 if TYPE_CHECKING:
     from google.api_core.retry import Retry
 

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

@@ -25,17 +25,19 @@ from typing import (
     cast,
 )
 
+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
 from google.api_core.exceptions import GoogleAPICallError
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
 from google.api_core.retry import Retry
 from google.cloud.translate_v2 import Client
 from google.cloud.translate_v3 import TranslationServiceClient
-
-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
+from google.cloud.translate_v3.types.translation_service import GlossaryInputConfig
 
 if TYPE_CHECKING:
+    from proto import Message
+
     from google.api_core.operation import Operation
     from google.cloud.translate_v3.services.translation_service import pagers
     from google.cloud.translate_v3.types import (
@@ -51,7 +53,7 @@ if TYPE_CHECKING:
         TransliterationConfig,
         automl_translation,
     )
-    from proto import Message
+    from google.cloud.translate_v3.types.translation_service import Glossary
 
 
 class WaitOperationNotDoneYetError(Exception):
@@ -915,3 +917,232 @@ class TranslateHook(GoogleBaseHook):
             retry=retry,
             metadata=metadata,
         )
+
+    def create_glossary(
+        self,
+        project_id: str,
+        location: str,
+        glossary_id: str,
+        input_config: GlossaryInputConfig | dict,
+        language_pair: Glossary.LanguageCodePair | dict | None = None,
+        language_codes_set: Glossary.LanguageCodesSet | MutableSequence[str] | None = None,
+        retry: Retry | _MethodDefault = DEFAULT,
+        timeout: float | None = None,
+        metadata: Sequence[tuple[str, str]] = (),
+    ) -> Operation:
+        """
+        Create the glossary resource from the input source file.
+
+        :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 glossary_id: User-specified id to built glossary resource name.
+        :param input_config: The input configuration of examples to built glossary from.
+            Total glossary must not exceed 10M Unicode codepoints.
+            The headers should not be included into the input file table, as languages specified with the
+            ``language_pair`` or ``language_codes_set`` params.
+        :param language_pair: Pair of language codes to be used for glossary creation.
+            Used to built unidirectional glossary. If specified, the ``language_codes_set`` should be empty.
+        :param language_codes_set: Set of language codes to create the equivalent term sets glossary.
+            Meant multiple languages mapping. If specified, the ``language_pair`` should be empty.
+        :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 glossary creation results.
+        """
+        client = self.get_client()
+        parent = f"projects/{project_id}/locations/{location}"
+        name = f"projects/{project_id}/locations/{location}/glossaries/{glossary_id}"
+
+        result = client.create_glossary(
+            request={
+                "parent": parent,
+                "glossary": {
+                    "name": name,
+                    "input_config": input_config,
+                    "language_pair": language_pair,
+                    "language_codes_set": language_codes_set,
+                },
+            },
+            retry=retry,
+            timeout=timeout,
+            metadata=metadata,
+        )
+        return result
+
+    def get_glossary(
+        self,
+        project_id: str,
+        location: str,
+        glossary_id: str,
+        retry: Retry | _MethodDefault = DEFAULT,
+        timeout: float | None = None,
+        metadata: Sequence[tuple[str, str]] = (),
+    ) -> Glossary:
+        """
+        Fetch glossary item data by the given id.
+
+        The glossary_id is a substring of glossary name, following the format:
+        ``projects/{project-number-or-id}/locations/{location-id}/glossaries/{glossary-id}``
+
+        :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 glossary_id: User-specified id to built glossary resource name.
+        :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: Fetched glossary item.
+        """
+        client = self.get_client()
+        name = f"projects/{project_id}/locations/{location}/glossaries/{glossary_id}"
+        result = client.get_glossary(
+            name=name,
+            retry=retry,
+            timeout=timeout,
+            metadata=metadata,
+        )
+        if not result:
+            raise AirflowException(f"Fail to get glossary {name}! Please check if it exists.")
+        return result
+
+    def update_glossary(
+        self,
+        glossary: Glossary,
+        new_display_name: str | None = None,
+        new_input_config: GlossaryInputConfig | dict | None = None,
+        retry: Retry | _MethodDefault = DEFAULT,
+        timeout: float | None = None,
+        metadata: Sequence[tuple[str, str]] = (),
+    ) -> Operation:
+        """
+        Update glossary item with values provided.
+
+        Only ``display_name`` and ``input_config`` fields are allowed for update.
+
+        :param glossary: Glossary item to update.
+        :param new_display_name: New value of the ``display_name`` to be updated.
+        :param new_input_config: New value of the ``input_config`` to be updated.
+        :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` with glossary update results.
+        """
+        client = self.get_client()
+        updated_fields = []
+        if new_display_name:
+            glossary.display_name = new_display_name
+            updated_fields.append("display_name")
+        if new_input_config is not None:
+            if isinstance(new_input_config, dict):
+                new_input_config = GlossaryInputConfig(**new_input_config)
+            glossary.input_config = new_input_config
+            updated_fields.append("input_config")
+        result = client.update_glossary(
+            request={"glossary": glossary, "update_mask": {"paths": updated_fields}},
+            retry=retry,
+            timeout=timeout,
+            metadata=metadata,
+        )
+        return result
+
+    def list_glossaries(
+        self,
+        project_id: str,
+        location: str,
+        page_size: int | None = None,
+        page_token: str | None = None,
+        filter_str: str | None = None,
+        retry: Retry | _MethodDefault = DEFAULT,
+        timeout: float | None = None,
+        metadata: Sequence[tuple[str, str]] = (),
+    ) -> pagers.ListGlossariesPager:
+        """
+        Get the list of glossaries available.
+
+        :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 page_size: Page size requested, if not set server use appropriate default.
+        :param page_token: A token identifying a page of results the server should return.
+            The first page is returned if ``page_token`` is empty or missing.
+        :param filter_str: Filter specifying constraints of a list operation. Specify the constraint by the
+            format of "key=value", where key must be ``src`` or ``tgt``, and the value must be a valid
+            language code.
+            For multiple restrictions, concatenate them by "AND" (uppercase only), such as:
+            ``src=en-US AND tgt=zh-CN``. Notice that the exact match is used here, which means using 'en-US'
+            and 'en' can lead to different results, which depends on the language code you used when you
+            create the glossary.
+            For the unidirectional glossaries, the ``src`` and ``tgt`` add restrictions
+            on the source and target language code separately.
+            For the equivalent term set glossaries, the ``src`` and/or ``tgt`` add restrictions on the term set.
+            For example: ``src=en-US AND tgt=zh-CN`` will only pick the unidirectional glossaries which exactly
+            match the source language code as ``en-US`` and the target language code ``zh-CN``, but all
+            equivalent term set glossaries which contain ``en-US`` and ``zh-CN`` in their language set will
+            be picked.
+            If missing, no filtering is performed.
+        :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: Glossaries list pager object.
+        """
+        client = self.get_client()
+        parent = f"projects/{project_id}/locations/{location}"
+        result = client.list_glossaries(
+            request={
+                "parent": parent,
+                "page_size": page_size,
+                "page_token": page_token,
+                "filter": filter_str,
+            },
+            retry=retry,
+            timeout=timeout,
+            metadata=metadata,
+        )
+        return result
+
+    def delete_glossary(
+        self,
+        project_id: str,
+        location: str,
+        glossary_id: str,
+        retry: Retry | _MethodDefault = DEFAULT,
+        timeout: float | None = None,
+        metadata: Sequence[tuple[str, str]] = (),
+    ) -> Operation:
+        """
+        Delete the glossary item by the given id.
+
+        :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 glossary_id: Glossary id to be deleted.
+        :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` with glossary deletion results.
+        """
+        client = self.get_client()
+        name = f"projects/{project_id}/locations/{location}/glossaries/{glossary_id}"
+        result = client.delete_glossary(
+            name=name,
+            retry=retry,
+            timeout=timeout,
+            metadata=metadata,
+        )
+        return result

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

@@ -23,6 +23,9 @@ import warnings
 from collections.abc import Sequence
 from typing import TYPE_CHECKING
 
+from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
+from airflow.providers.google.common.deprecated import deprecated
+from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
 from google.api_core.client_options import ClientOptions
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
 from google.cloud.aiplatform import (
@@ -36,10 +39,6 @@ from google.cloud.aiplatform import (
 )
 from google.cloud.aiplatform_v1 import JobServiceClient, PipelineServiceClient
 
-from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
-from airflow.providers.google.common.deprecated import deprecated
-from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
-
 if TYPE_CHECKING:
     from google.api_core.operation import Operation
     from google.api_core.retry import Retry
@@ -561,6 +560,7 @@ class AutoMLHook(GoogleBaseHook):
         model_version_description: str | None = None,
         window_stride_length: int | None = None,
         window_max_count: int | None = None,
+        holiday_regions: list[str] | None = None,
     ) -> tuple[models.Model | None, str]:
         """
         Create an AutoML Forecasting Training Job.
@@ -717,6 +717,10 @@ class AutoMLHook(GoogleBaseHook):
             ``window_stride_length`` rows will be used to generate a sliding window.
         :param window_max_count: Optional. Number of rows that should be used to generate input examples. If the
             total row count is larger than this number, the input data will be randomly sampled to hit the count.
+        :param holiday_regions: Optional. You can select one or more geographical
+            regions to enable holiday effect modeling. During training, Vertex AI
+            creates holiday categorical features within the model based on the date
+            from TIME_COLUMN and the specified geographical regions.
         """
         if column_transformations:
             warnings.warn(
@@ -774,6 +778,7 @@ class AutoMLHook(GoogleBaseHook):
             model_version_description=model_version_description,
             window_stride_length=window_stride_length,
             window_max_count=window_max_count,
+            holiday_regions=holiday_regions,
         )
         training_id = self.extract_training_id(self._job.resource_name)
         if model:

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

@@ -23,15 +23,14 @@ import asyncio
 from collections.abc import Sequence
 from typing import TYPE_CHECKING
 
+from airflow.exceptions import AirflowException
+from airflow.providers.google.common.consts import CLIENT_INFO
+from airflow.providers.google.common.hooks.base_google import GoogleBaseAsyncHook, GoogleBaseHook
 from google.api_core.client_options import ClientOptions
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
 from google.cloud.aiplatform import BatchPredictionJob, Model, explain
 from google.cloud.aiplatform_v1 import JobServiceAsyncClient, JobServiceClient, JobState, types
 
-from airflow.exceptions import AirflowException
-from airflow.providers.google.common.consts import CLIENT_INFO
-from airflow.providers.google.common.hooks.base_google import GoogleBaseAsyncHook, GoogleBaseHook
-
 if TYPE_CHECKING:
     from google.api_core.operation import Operation
     from google.api_core.retry import AsyncRetry, Retry

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

@@ -23,6 +23,10 @@ import asyncio
 from collections.abc import Sequence
 from typing import TYPE_CHECKING, Any
 
+from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
+from airflow.providers.google.common.consts import CLIENT_INFO
+from airflow.providers.google.common.deprecated import deprecated
+from airflow.providers.google.common.hooks.base_google import GoogleBaseAsyncHook, GoogleBaseHook
 from google.api_core.client_options import ClientOptions
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
 from google.cloud.aiplatform import (
@@ -42,11 +46,6 @@ from google.cloud.aiplatform_v1 import (
     types,
 )
 
-from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
-from airflow.providers.google.common.consts import CLIENT_INFO
-from airflow.providers.google.common.deprecated import deprecated
-from airflow.providers.google.common.hooks.base_google import GoogleBaseAsyncHook, GoogleBaseHook
-
 if TYPE_CHECKING:
     from google.api_core.operation import Operation
     from google.api_core.retry import AsyncRetry, Retry

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

@@ -22,13 +22,12 @@ from __future__ import annotations
 from collections.abc import Sequence
 from typing import TYPE_CHECKING
 
-from google.api_core.client_options import ClientOptions
-from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
-from google.cloud.aiplatform_v1 import DatasetServiceClient
-
 from airflow.exceptions import AirflowException
 from airflow.providers.google.common.consts import CLIENT_INFO
 from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
+from google.api_core.client_options import ClientOptions
+from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
+from google.cloud.aiplatform_v1 import DatasetServiceClient
 
 if TYPE_CHECKING:
     from google.api_core.operation import Operation

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

@@ -22,13 +22,12 @@ from __future__ import annotations
 from collections.abc import Sequence
 from typing import TYPE_CHECKING
 
+from airflow.exceptions import AirflowException
+from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
 from google.api_core.client_options import ClientOptions
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
 from google.cloud.aiplatform_v1 import EndpointServiceClient
 
-from airflow.exceptions import AirflowException
-from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
-
 if TYPE_CHECKING:
     from google.api_core.operation import Operation
     from google.api_core.retry import Retry

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

@@ -18,15 +18,14 @@
 
 from __future__ import annotations
 
+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
 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):
     """