apache-airflow-providers-google 14.0.0rc1__py3-none-any.whl → 14.1.0rc1__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- airflow/providers/google/3rd-party-licenses/LICENSES.txt +14 -0
- airflow/providers/google/3rd-party-licenses/NOTICE +15 -0
- airflow/providers/google/__init__.py +1 -1
- airflow/providers/google/_vendor/__init__.py +0 -0
- airflow/providers/google/_vendor/json_merge_patch.py +91 -0
- airflow/providers/google/ads/hooks/ads.py +12 -11
- airflow/providers/google/cloud/_internal_client/secret_manager_client.py +3 -2
- airflow/providers/google/cloud/hooks/alloy_db.py +2 -3
- airflow/providers/google/cloud/hooks/automl.py +7 -13
- airflow/providers/google/cloud/hooks/bigquery.py +21 -22
- airflow/providers/google/cloud/hooks/bigquery_dts.py +8 -8
- airflow/providers/google/cloud/hooks/bigtable.py +3 -2
- airflow/providers/google/cloud/hooks/cloud_batch.py +4 -3
- airflow/providers/google/cloud/hooks/cloud_build.py +7 -13
- airflow/providers/google/cloud/hooks/cloud_composer.py +7 -12
- airflow/providers/google/cloud/hooks/cloud_memorystore.py +4 -3
- airflow/providers/google/cloud/hooks/cloud_run.py +4 -3
- airflow/providers/google/cloud/hooks/cloud_sql.py +1 -1
- airflow/providers/google/cloud/hooks/cloud_storage_transfer_service.py +8 -9
- airflow/providers/google/cloud/hooks/compute.py +3 -3
- airflow/providers/google/cloud/hooks/datacatalog.py +4 -3
- airflow/providers/google/cloud/hooks/dataflow.py +12 -12
- airflow/providers/google/cloud/hooks/dataform.py +3 -2
- airflow/providers/google/cloud/hooks/datafusion.py +2 -2
- airflow/providers/google/cloud/hooks/dataplex.py +11 -10
- airflow/providers/google/cloud/hooks/dataproc.py +5 -4
- airflow/providers/google/cloud/hooks/dataproc_metastore.py +4 -3
- airflow/providers/google/cloud/hooks/dlp.py +4 -3
- airflow/providers/google/cloud/hooks/gcs.py +19 -12
- airflow/providers/google/cloud/hooks/kms.py +3 -2
- airflow/providers/google/cloud/hooks/kubernetes_engine.py +20 -14
- airflow/providers/google/cloud/hooks/life_sciences.py +1 -1
- airflow/providers/google/cloud/hooks/managed_kafka.py +227 -3
- airflow/providers/google/cloud/hooks/natural_language.py +3 -2
- airflow/providers/google/cloud/hooks/os_login.py +3 -2
- airflow/providers/google/cloud/hooks/pubsub.py +6 -6
- airflow/providers/google/cloud/hooks/secret_manager.py +3 -2
- airflow/providers/google/cloud/hooks/spanner.py +2 -2
- airflow/providers/google/cloud/hooks/speech_to_text.py +3 -2
- airflow/providers/google/cloud/hooks/stackdriver.py +4 -4
- airflow/providers/google/cloud/hooks/tasks.py +4 -3
- airflow/providers/google/cloud/hooks/text_to_speech.py +3 -2
- airflow/providers/google/cloud/hooks/translate.py +6 -15
- airflow/providers/google/cloud/hooks/vertex_ai/auto_ml.py +6 -12
- airflow/providers/google/cloud/hooks/vertex_ai/batch_prediction_job.py +6 -12
- airflow/providers/google/cloud/hooks/vertex_ai/custom_job.py +7 -13
- airflow/providers/google/cloud/hooks/vertex_ai/dataset.py +5 -12
- airflow/providers/google/cloud/hooks/vertex_ai/endpoint_service.py +4 -11
- airflow/providers/google/cloud/hooks/vertex_ai/feature_store.py +4 -3
- airflow/providers/google/cloud/hooks/vertex_ai/hyperparameter_tuning_job.py +6 -12
- airflow/providers/google/cloud/hooks/vertex_ai/model_service.py +6 -11
- airflow/providers/google/cloud/hooks/vertex_ai/pipeline_job.py +6 -12
- airflow/providers/google/cloud/hooks/vertex_ai/prediction_service.py +3 -2
- airflow/providers/google/cloud/hooks/video_intelligence.py +3 -2
- airflow/providers/google/cloud/hooks/vision.py +4 -3
- airflow/providers/google/cloud/hooks/workflows.py +3 -2
- airflow/providers/google/cloud/links/base.py +7 -1
- airflow/providers/google/cloud/links/datafusion.py +8 -2
- airflow/providers/google/cloud/links/dataproc.py +8 -1
- airflow/providers/google/cloud/links/kubernetes_engine.py +2 -1
- airflow/providers/google/cloud/links/managed_kafka.py +30 -0
- airflow/providers/google/cloud/log/gcs_task_handler.py +15 -7
- airflow/providers/google/cloud/log/stackdriver_task_handler.py +8 -6
- airflow/providers/google/cloud/openlineage/utils.py +4 -2
- airflow/providers/google/cloud/operators/alloy_db.py +6 -5
- airflow/providers/google/cloud/operators/automl.py +11 -9
- airflow/providers/google/cloud/operators/bigquery.py +8 -6
- airflow/providers/google/cloud/operators/bigquery_dts.py +10 -8
- airflow/providers/google/cloud/operators/bigtable.py +3 -1
- airflow/providers/google/cloud/operators/cloud_base.py +2 -1
- airflow/providers/google/cloud/operators/cloud_batch.py +4 -2
- airflow/providers/google/cloud/operators/cloud_build.py +5 -3
- airflow/providers/google/cloud/operators/cloud_composer.py +7 -5
- airflow/providers/google/cloud/operators/cloud_memorystore.py +6 -4
- airflow/providers/google/cloud/operators/cloud_run.py +5 -3
- airflow/providers/google/cloud/operators/compute.py +5 -4
- airflow/providers/google/cloud/operators/datacatalog.py +11 -9
- airflow/providers/google/cloud/operators/dataform.py +5 -3
- airflow/providers/google/cloud/operators/datafusion.py +1 -1
- airflow/providers/google/cloud/operators/dataplex.py +20 -18
- airflow/providers/google/cloud/operators/dataproc.py +34 -5
- airflow/providers/google/cloud/operators/dataproc_metastore.py +18 -7
- airflow/providers/google/cloud/operators/dlp.py +19 -17
- airflow/providers/google/cloud/operators/gcs.py +5 -4
- airflow/providers/google/cloud/operators/kubernetes_engine.py +54 -5
- airflow/providers/google/cloud/operators/managed_kafka.py +271 -4
- airflow/providers/google/cloud/operators/natural_language.py +5 -3
- airflow/providers/google/cloud/operators/pubsub.py +10 -8
- airflow/providers/google/cloud/operators/speech_to_text.py +5 -3
- airflow/providers/google/cloud/operators/stackdriver.py +5 -3
- airflow/providers/google/cloud/operators/tasks.py +6 -4
- airflow/providers/google/cloud/operators/text_to_speech.py +4 -2
- airflow/providers/google/cloud/operators/translate.py +5 -3
- airflow/providers/google/cloud/operators/translate_speech.py +4 -2
- airflow/providers/google/cloud/operators/vertex_ai/auto_ml.py +8 -6
- airflow/providers/google/cloud/operators/vertex_ai/batch_prediction_job.py +6 -4
- airflow/providers/google/cloud/operators/vertex_ai/custom_job.py +8 -6
- airflow/providers/google/cloud/operators/vertex_ai/dataset.py +6 -4
- airflow/providers/google/cloud/operators/vertex_ai/endpoint_service.py +6 -4
- airflow/providers/google/cloud/operators/vertex_ai/hyperparameter_tuning_job.py +6 -4
- airflow/providers/google/cloud/operators/vertex_ai/model_service.py +6 -4
- airflow/providers/google/cloud/operators/vertex_ai/pipeline_job.py +6 -4
- airflow/providers/google/cloud/operators/video_intelligence.py +5 -3
- airflow/providers/google/cloud/operators/vision.py +6 -4
- airflow/providers/google/cloud/operators/workflows.py +7 -5
- airflow/providers/google/cloud/secrets/secret_manager.py +2 -1
- airflow/providers/google/cloud/sensors/bigquery_dts.py +5 -3
- airflow/providers/google/cloud/sensors/bigtable.py +3 -2
- airflow/providers/google/cloud/sensors/cloud_composer.py +1 -1
- airflow/providers/google/cloud/sensors/dataplex.py +6 -4
- airflow/providers/google/cloud/sensors/dataproc.py +3 -2
- airflow/providers/google/cloud/sensors/dataproc_metastore.py +2 -1
- airflow/providers/google/cloud/sensors/gcs.py +4 -2
- airflow/providers/google/cloud/sensors/pubsub.py +3 -2
- airflow/providers/google/cloud/sensors/workflows.py +5 -3
- airflow/providers/google/cloud/transfers/bigquery_to_gcs.py +5 -3
- airflow/providers/google/cloud/transfers/gcs_to_bigquery.py +12 -10
- airflow/providers/google/cloud/triggers/bigquery_dts.py +2 -1
- airflow/providers/google/cloud/triggers/cloud_batch.py +2 -1
- airflow/providers/google/cloud/triggers/cloud_build.py +2 -1
- airflow/providers/google/cloud/triggers/cloud_composer.py +1 -1
- airflow/providers/google/cloud/triggers/cloud_storage_transfer_service.py +3 -2
- airflow/providers/google/cloud/triggers/dataflow.py +3 -2
- airflow/providers/google/cloud/triggers/dataplex.py +2 -1
- airflow/providers/google/cloud/triggers/dataproc.py +3 -2
- airflow/providers/google/cloud/triggers/kubernetes_engine.py +1 -1
- airflow/providers/google/cloud/triggers/pubsub.py +2 -1
- airflow/providers/google/cloud/triggers/vertex_ai.py +8 -7
- airflow/providers/google/cloud/utils/credentials_provider.py +4 -3
- airflow/providers/google/cloud/utils/external_token_supplier.py +0 -1
- airflow/providers/google/cloud/{example_dags/__init__.py → utils/validators.py} +27 -0
- airflow/providers/google/common/auth_backend/google_openid.py +14 -5
- airflow/providers/google/common/consts.py +2 -1
- airflow/providers/google/common/hooks/base_google.py +7 -8
- airflow/providers/google/common/hooks/operation_helpers.py +79 -0
- airflow/providers/google/get_provider_info.py +11 -6
- airflow/providers/google/marketing_platform/hooks/analytics_admin.py +3 -2
- airflow/providers/google/marketing_platform/hooks/search_ads.py +1 -1
- airflow/providers/google/marketing_platform/links/analytics_admin.py +10 -2
- airflow/providers/google/marketing_platform/operators/analytics_admin.py +7 -5
- {apache_airflow_providers_google-14.0.0rc1.dist-info → apache_airflow_providers_google-14.1.0rc1.dist-info}/METADATA +32 -35
- {apache_airflow_providers_google-14.0.0rc1.dist-info → apache_airflow_providers_google-14.1.0rc1.dist-info}/RECORD +144 -140
- {apache_airflow_providers_google-14.0.0rc1.dist-info → apache_airflow_providers_google-14.1.0rc1.dist-info}/WHEEL +1 -1
- airflow/providers/google/cloud/example_dags/example_cloud_task.py +0 -54
- {apache_airflow_providers_google-14.0.0rc1.dist-info → apache_airflow_providers_google-14.1.0rc1.dist-info}/entry_points.txt +0 -0
@@ -31,18 +31,6 @@ from collections.abc import Generator, Sequence
|
|
31
31
|
from copy import deepcopy
|
32
32
|
from typing import TYPE_CHECKING, Any, Callable, TypeVar, cast
|
33
33
|
|
34
|
-
from googleapiclient.discovery import Resource, build
|
35
|
-
|
36
|
-
from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
|
37
|
-
from airflow.providers.apache.beam.hooks.beam import BeamHook, BeamRunnerType, beam_options_to_args
|
38
|
-
from airflow.providers.google.common.deprecated import deprecated
|
39
|
-
from airflow.providers.google.common.hooks.base_google import (
|
40
|
-
PROVIDE_PROJECT_ID,
|
41
|
-
GoogleBaseAsyncHook,
|
42
|
-
GoogleBaseHook,
|
43
|
-
)
|
44
|
-
from airflow.utils.log.logging_mixin import LoggingMixin
|
45
|
-
from airflow.utils.timeout import timeout
|
46
34
|
from google.cloud.dataflow_v1beta3 import (
|
47
35
|
GetJobRequest,
|
48
36
|
Job,
|
@@ -59,6 +47,18 @@ from google.cloud.dataflow_v1beta3.types import (
|
|
59
47
|
JobMetrics,
|
60
48
|
)
|
61
49
|
from google.cloud.dataflow_v1beta3.types.jobs import ListJobsRequest
|
50
|
+
from googleapiclient.discovery import Resource, build
|
51
|
+
|
52
|
+
from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
|
53
|
+
from airflow.providers.apache.beam.hooks.beam import BeamHook, BeamRunnerType, beam_options_to_args
|
54
|
+
from airflow.providers.google.common.deprecated import deprecated
|
55
|
+
from airflow.providers.google.common.hooks.base_google import (
|
56
|
+
PROVIDE_PROJECT_ID,
|
57
|
+
GoogleBaseAsyncHook,
|
58
|
+
GoogleBaseHook,
|
59
|
+
)
|
60
|
+
from airflow.utils.log.logging_mixin import LoggingMixin
|
61
|
+
from airflow.utils.timeout import timeout
|
62
62
|
|
63
63
|
if TYPE_CHECKING:
|
64
64
|
from google.cloud.dataflow_v1beta3.services.jobs_v1_beta3.pagers import ListJobsAsyncPager
|
@@ -20,8 +20,6 @@ import time
|
|
20
20
|
from collections.abc import Sequence
|
21
21
|
from typing import TYPE_CHECKING
|
22
22
|
|
23
|
-
from airflow.exceptions import AirflowException
|
24
|
-
from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
|
25
23
|
from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
|
26
24
|
from google.cloud.dataform_v1beta1 import DataformClient
|
27
25
|
from google.cloud.dataform_v1beta1.types import (
|
@@ -33,6 +31,9 @@ from google.cloud.dataform_v1beta1.types import (
|
|
33
31
|
WriteFileResponse,
|
34
32
|
)
|
35
33
|
|
34
|
+
from airflow.exceptions import AirflowException
|
35
|
+
from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
|
36
|
+
|
36
37
|
if TYPE_CHECKING:
|
37
38
|
from google.api_core.retry import Retry
|
38
39
|
from google.cloud.dataform_v1beta1.services.dataform.pagers import QueryWorkflowInvocationActionsPager
|
@@ -26,11 +26,12 @@ from collections.abc import Sequence
|
|
26
26
|
from typing import Any
|
27
27
|
from urllib.parse import quote, urlencode, urljoin
|
28
28
|
|
29
|
+
import google.auth
|
29
30
|
from aiohttp import ClientSession
|
30
31
|
from gcloud.aio.auth import AioSession, Token
|
32
|
+
from google.api_core.retry import exponential_sleep_generator
|
31
33
|
from googleapiclient.discovery import Resource, build
|
32
34
|
|
33
|
-
import google.auth
|
34
35
|
from airflow.exceptions import AirflowException, AirflowNotFoundException
|
35
36
|
from airflow.providers.google.cloud.utils.datafusion import DataFusionPipelineType
|
36
37
|
from airflow.providers.google.common.hooks.base_google import (
|
@@ -38,7 +39,6 @@ from airflow.providers.google.common.hooks.base_google import (
|
|
38
39
|
GoogleBaseAsyncHook,
|
39
40
|
GoogleBaseHook,
|
40
41
|
)
|
41
|
-
from google.api_core.retry import exponential_sleep_generator
|
42
42
|
|
43
43
|
Operation = dict[str, Any]
|
44
44
|
|
@@ -23,13 +23,6 @@ from collections.abc import MutableSequence, Sequence
|
|
23
23
|
from copy import deepcopy
|
24
24
|
from typing import TYPE_CHECKING, Any
|
25
25
|
|
26
|
-
from airflow.exceptions import AirflowException
|
27
|
-
from airflow.providers.google.common.consts import CLIENT_INFO
|
28
|
-
from airflow.providers.google.common.hooks.base_google import (
|
29
|
-
PROVIDE_PROJECT_ID,
|
30
|
-
GoogleBaseAsyncHook,
|
31
|
-
GoogleBaseHook,
|
32
|
-
)
|
33
26
|
from google.api_core.client_options import ClientOptions
|
34
27
|
from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
|
35
28
|
from google.cloud.dataplex_v1 import (
|
@@ -53,9 +46,16 @@ from google.cloud.dataplex_v1.types import (
|
|
53
46
|
)
|
54
47
|
from google.protobuf.field_mask_pb2 import FieldMask
|
55
48
|
|
56
|
-
|
57
|
-
|
49
|
+
from airflow.exceptions import AirflowException
|
50
|
+
from airflow.providers.google.common.consts import CLIENT_INFO
|
51
|
+
from airflow.providers.google.common.hooks.base_google import (
|
52
|
+
PROVIDE_PROJECT_ID,
|
53
|
+
GoogleBaseAsyncHook,
|
54
|
+
GoogleBaseHook,
|
55
|
+
)
|
56
|
+
from airflow.providers.google.common.hooks.operation_helpers import OperationHelper
|
58
57
|
|
58
|
+
if TYPE_CHECKING:
|
59
59
|
from google.api_core.operation import Operation
|
60
60
|
from google.api_core.retry import Retry
|
61
61
|
from google.api_core.retry_async import AsyncRetry
|
@@ -66,6 +66,7 @@ if TYPE_CHECKING:
|
|
66
66
|
ListEntryTypesPager,
|
67
67
|
SearchEntriesPager,
|
68
68
|
)
|
69
|
+
from googleapiclient.discovery import Resource
|
69
70
|
|
70
71
|
PATH_DATA_SCAN = "projects/{project_id}/locations/{region}/dataScans/{data_scan_id}"
|
71
72
|
|
@@ -78,7 +79,7 @@ class AirflowDataQualityScanResultTimeoutException(AirflowException):
|
|
78
79
|
"""Raised when no result found after specified amount of seconds."""
|
79
80
|
|
80
81
|
|
81
|
-
class DataplexHook(GoogleBaseHook):
|
82
|
+
class DataplexHook(GoogleBaseHook, OperationHelper):
|
82
83
|
"""
|
83
84
|
Hook for Google Dataplex.
|
84
85
|
|
@@ -26,10 +26,6 @@ import uuid
|
|
26
26
|
from collections.abc import MutableSequence, Sequence
|
27
27
|
from typing import TYPE_CHECKING, Any
|
28
28
|
|
29
|
-
from airflow.exceptions import AirflowException
|
30
|
-
from airflow.providers.google.common.consts import CLIENT_INFO
|
31
|
-
from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
|
32
|
-
from airflow.version import version as airflow_version
|
33
29
|
from google.api_core.client_options import ClientOptions
|
34
30
|
from google.api_core.exceptions import ServerError
|
35
31
|
from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
|
@@ -49,6 +45,11 @@ from google.cloud.dataproc_v1 import (
|
|
49
45
|
WorkflowTemplateServiceClient,
|
50
46
|
)
|
51
47
|
|
48
|
+
from airflow.exceptions import AirflowException
|
49
|
+
from airflow.providers.google.common.consts import CLIENT_INFO
|
50
|
+
from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
|
51
|
+
from airflow.version import version as airflow_version
|
52
|
+
|
52
53
|
if TYPE_CHECKING:
|
53
54
|
from google.api_core.operation import Operation
|
54
55
|
from google.api_core.operation_async import AsyncOperation
|
@@ -22,13 +22,14 @@ from __future__ import annotations
|
|
22
22
|
from collections.abc import Sequence
|
23
23
|
from typing import TYPE_CHECKING, Any
|
24
24
|
|
25
|
-
from airflow.exceptions import AirflowException
|
26
|
-
from airflow.providers.google.common.consts import CLIENT_INFO
|
27
|
-
from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
|
28
25
|
from google.api_core.client_options import ClientOptions
|
29
26
|
from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
|
30
27
|
from google.cloud.metastore_v1 import DataprocMetastoreClient
|
31
28
|
|
29
|
+
from airflow.exceptions import AirflowException
|
30
|
+
from airflow.providers.google.common.consts import CLIENT_INFO
|
31
|
+
from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
|
32
|
+
|
32
33
|
if TYPE_CHECKING:
|
33
34
|
from google.api_core.operation import Operation
|
34
35
|
from google.api_core.retry import Retry
|
@@ -31,9 +31,6 @@ import time
|
|
31
31
|
from collections.abc import Sequence
|
32
32
|
from typing import TYPE_CHECKING
|
33
33
|
|
34
|
-
from airflow.exceptions import AirflowException
|
35
|
-
from airflow.providers.google.common.consts import CLIENT_INFO
|
36
|
-
from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID, GoogleBaseHook
|
37
34
|
from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
|
38
35
|
from google.cloud.dlp import DlpServiceClient
|
39
36
|
from google.cloud.dlp_v2.types import (
|
@@ -58,6 +55,10 @@ from google.cloud.dlp_v2.types import (
|
|
58
55
|
)
|
59
56
|
from google.protobuf.field_mask_pb2 import FieldMask
|
60
57
|
|
58
|
+
from airflow.exceptions import AirflowException
|
59
|
+
from airflow.providers.google.common.consts import CLIENT_INFO
|
60
|
+
from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID, GoogleBaseHook
|
61
|
+
|
61
62
|
if TYPE_CHECKING:
|
62
63
|
from google.api_core.retry import Retry
|
63
64
|
|
@@ -35,6 +35,12 @@ from typing import IO, TYPE_CHECKING, Any, Callable, TypeVar, cast, overload
|
|
35
35
|
from urllib.parse import urlsplit
|
36
36
|
|
37
37
|
from gcloud.aio.storage import Storage
|
38
|
+
from google.api_core.exceptions import GoogleAPICallError, NotFound
|
39
|
+
|
40
|
+
# not sure why but mypy complains on missing `storage` but it is clearly there and is importable
|
41
|
+
from google.cloud import storage # type: ignore[attr-defined]
|
42
|
+
from google.cloud.exceptions import GoogleCloudError
|
43
|
+
from google.cloud.storage.retry import DEFAULT_RETRY
|
38
44
|
from requests import Session
|
39
45
|
|
40
46
|
from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
|
@@ -49,18 +55,11 @@ from airflow.providers.google.common.hooks.base_google import (
|
|
49
55
|
from airflow.typing_compat import ParamSpec
|
50
56
|
from airflow.utils import timezone
|
51
57
|
from airflow.version import version
|
52
|
-
from google.api_core.exceptions import GoogleAPICallError, NotFound
|
53
|
-
|
54
|
-
# not sure why but mypy complains on missing `storage` but it is clearly there and is importable
|
55
|
-
from google.cloud import storage # type: ignore[attr-defined]
|
56
|
-
from google.cloud.exceptions import GoogleCloudError
|
57
|
-
from google.cloud.storage.retry import DEFAULT_RETRY
|
58
58
|
|
59
59
|
if TYPE_CHECKING:
|
60
60
|
from datetime import datetime
|
61
61
|
|
62
62
|
from aiohttp import ClientSession
|
63
|
-
|
64
63
|
from google.api_core.retry import Retry
|
65
64
|
from google.cloud.storage.blob import Blob
|
66
65
|
|
@@ -598,7 +597,13 @@ class GCSHook(GoogleBaseHook):
|
|
598
597
|
context=self, scheme="gs", asset_kwargs={"bucket": bucket.name, "key": blob.name}
|
599
598
|
)
|
600
599
|
|
601
|
-
def exists(
|
600
|
+
def exists(
|
601
|
+
self,
|
602
|
+
bucket_name: str,
|
603
|
+
object_name: str,
|
604
|
+
retry: Retry = DEFAULT_RETRY,
|
605
|
+
user_project: str | None = None,
|
606
|
+
) -> bool:
|
602
607
|
"""
|
603
608
|
Check for the existence of a file in Google Cloud Storage.
|
604
609
|
|
@@ -606,9 +611,11 @@ class GCSHook(GoogleBaseHook):
|
|
606
611
|
:param object_name: The name of the blob_name to check in the Google cloud
|
607
612
|
storage bucket.
|
608
613
|
:param retry: (Optional) How to retry the RPC
|
614
|
+
:param user_project: The identifier of the Google Cloud project to bill for the request.
|
615
|
+
Required for Requester Pays buckets.
|
609
616
|
"""
|
610
617
|
client = self.get_conn()
|
611
|
-
bucket = client.bucket(bucket_name)
|
618
|
+
bucket = client.bucket(bucket_name, user_project=user_project)
|
612
619
|
blob = bucket.blob(blob_name=object_name)
|
613
620
|
return blob.exists(retry=retry)
|
614
621
|
|
@@ -625,7 +632,7 @@ class GCSHook(GoogleBaseHook):
|
|
625
632
|
|
626
633
|
def is_updated_after(self, bucket_name: str, object_name: str, ts: datetime) -> bool:
|
627
634
|
"""
|
628
|
-
Check if
|
635
|
+
Check if a blob_name is updated in Google Cloud Storage.
|
629
636
|
|
630
637
|
:param bucket_name: The Google Cloud Storage bucket where the object is.
|
631
638
|
:param object_name: The name of the object to check in the Google cloud
|
@@ -645,7 +652,7 @@ class GCSHook(GoogleBaseHook):
|
|
645
652
|
self, bucket_name: str, object_name: str, min_ts: datetime, max_ts: datetime
|
646
653
|
) -> bool:
|
647
654
|
"""
|
648
|
-
Check if
|
655
|
+
Check if a blob_name is updated in Google Cloud Storage.
|
649
656
|
|
650
657
|
:param bucket_name: The Google Cloud Storage bucket where the object is.
|
651
658
|
:param object_name: The name of the object to check in the Google cloud
|
@@ -666,7 +673,7 @@ class GCSHook(GoogleBaseHook):
|
|
666
673
|
|
667
674
|
def is_updated_before(self, bucket_name: str, object_name: str, ts: datetime) -> bool:
|
668
675
|
"""
|
669
|
-
Check if
|
676
|
+
Check if a blob_name is updated before given time in Google Cloud Storage.
|
670
677
|
|
671
678
|
:param bucket_name: The Google Cloud Storage bucket where the object is.
|
672
679
|
:param object_name: The name of the object to check in the Google cloud
|
@@ -23,11 +23,12 @@ import base64
|
|
23
23
|
from collections.abc import Sequence
|
24
24
|
from typing import TYPE_CHECKING
|
25
25
|
|
26
|
-
from airflow.providers.google.common.consts import CLIENT_INFO
|
27
|
-
from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
|
28
26
|
from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
|
29
27
|
from google.cloud.kms_v1 import KeyManagementServiceClient
|
30
28
|
|
29
|
+
from airflow.providers.google.common.consts import CLIENT_INFO
|
30
|
+
from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
|
31
|
+
|
31
32
|
if TYPE_CHECKING:
|
32
33
|
from google.api_core.retry import Retry
|
33
34
|
|
@@ -25,6 +25,14 @@ import time
|
|
25
25
|
from collections.abc import Sequence
|
26
26
|
from typing import TYPE_CHECKING, Any
|
27
27
|
|
28
|
+
from google.api_core.exceptions import NotFound
|
29
|
+
from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
|
30
|
+
from google.auth.transport import requests as google_requests
|
31
|
+
|
32
|
+
# not sure why but mypy complains on missing `container_v1` but it is clearly there and is importable
|
33
|
+
from google.cloud import exceptions # type: ignore[attr-defined]
|
34
|
+
from google.cloud.container_v1 import ClusterManagerAsyncClient, ClusterManagerClient
|
35
|
+
from google.cloud.container_v1.types import Cluster, Operation
|
28
36
|
from kubernetes import client
|
29
37
|
from kubernetes_asyncio import client as async_client
|
30
38
|
from kubernetes_asyncio.config.kube_config import FileOrData
|
@@ -39,14 +47,6 @@ from airflow.providers.google.common.hooks.base_google import (
|
|
39
47
|
GoogleBaseAsyncHook,
|
40
48
|
GoogleBaseHook,
|
41
49
|
)
|
42
|
-
from google.api_core.exceptions import NotFound
|
43
|
-
from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
|
44
|
-
from google.auth.transport import requests as google_requests
|
45
|
-
|
46
|
-
# not sure why but mypy complains on missing `container_v1` but it is clearly there and is importable
|
47
|
-
from google.cloud import exceptions # type: ignore[attr-defined]
|
48
|
-
from google.cloud.container_v1 import ClusterManagerAsyncClient, ClusterManagerClient
|
49
|
-
from google.cloud.container_v1.types import Cluster, Operation
|
50
50
|
|
51
51
|
if TYPE_CHECKING:
|
52
52
|
import google.auth.credentials
|
@@ -64,11 +64,13 @@ class GKEClusterConnection:
|
|
64
64
|
ssl_ca_cert: str,
|
65
65
|
credentials: google.auth.credentials.Credentials,
|
66
66
|
enable_tcp_keepalive: bool = False,
|
67
|
+
use_dns_endpoint: bool = False,
|
67
68
|
):
|
68
69
|
self._cluster_url = cluster_url
|
69
70
|
self._ssl_ca_cert = ssl_ca_cert
|
70
71
|
self._credentials = credentials
|
71
72
|
self.enable_tcp_keepalive = enable_tcp_keepalive
|
73
|
+
self.use_dns_endpoint = use_dns_endpoint
|
72
74
|
|
73
75
|
def get_conn(self) -> client.ApiClient:
|
74
76
|
configuration = self._get_config()
|
@@ -86,12 +88,13 @@ class GKEClusterConnection:
|
|
86
88
|
api_key_prefix={"authorization": "Bearer"},
|
87
89
|
api_key={"authorization": self._get_token(self._credentials)},
|
88
90
|
)
|
89
|
-
|
90
|
-
|
91
|
-
|
92
|
-
|
93
|
-
|
94
|
-
|
91
|
+
if not self.use_dns_endpoint:
|
92
|
+
configuration.ssl_ca_cert = FileOrData(
|
93
|
+
{
|
94
|
+
"certificate-authority-data": self._ssl_ca_cert,
|
95
|
+
},
|
96
|
+
file_key_name="certificate-authority",
|
97
|
+
).as_file()
|
95
98
|
return configuration
|
96
99
|
|
97
100
|
@staticmethod
|
@@ -417,6 +420,7 @@ class GKEKubernetesHook(GoogleBaseHook, KubernetesHook):
|
|
417
420
|
cluster_url: str,
|
418
421
|
ssl_ca_cert: str,
|
419
422
|
enable_tcp_keepalive: bool = False,
|
423
|
+
use_dns_endpoint: bool = False,
|
420
424
|
*args,
|
421
425
|
**kwargs,
|
422
426
|
):
|
@@ -424,6 +428,7 @@ class GKEKubernetesHook(GoogleBaseHook, KubernetesHook):
|
|
424
428
|
self._cluster_url = cluster_url
|
425
429
|
self._ssl_ca_cert = ssl_ca_cert
|
426
430
|
self.enable_tcp_keepalive = enable_tcp_keepalive
|
431
|
+
self.use_dns_endpoint = use_dns_endpoint
|
427
432
|
|
428
433
|
def get_conn(self) -> client.ApiClient:
|
429
434
|
return GKEClusterConnection(
|
@@ -431,6 +436,7 @@ class GKEKubernetesHook(GoogleBaseHook, KubernetesHook):
|
|
431
436
|
ssl_ca_cert=self._ssl_ca_cert,
|
432
437
|
credentials=self.get_credentials(),
|
433
438
|
enable_tcp_keepalive=self.enable_tcp_keepalive,
|
439
|
+
use_dns_endpoint=self.use_dns_endpoint,
|
434
440
|
).get_conn()
|
435
441
|
|
436
442
|
def apply_from_yaml_file(
|
@@ -22,9 +22,9 @@ from __future__ import annotations
|
|
22
22
|
import time
|
23
23
|
from collections.abc import Sequence
|
24
24
|
|
25
|
+
import google.api_core.path_template
|
25
26
|
from googleapiclient.discovery import build
|
26
27
|
|
27
|
-
import google.api_core.path_template
|
28
28
|
from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
|
29
29
|
from airflow.providers.google.common.deprecated import deprecated
|
30
30
|
from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
|
@@ -19,23 +19,81 @@
|
|
19
19
|
|
20
20
|
from __future__ import annotations
|
21
21
|
|
22
|
+
import base64
|
23
|
+
import datetime
|
24
|
+
import json
|
25
|
+
import time
|
22
26
|
from collections.abc import Sequence
|
23
27
|
from copy import deepcopy
|
24
28
|
from typing import TYPE_CHECKING
|
25
29
|
|
30
|
+
from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
|
31
|
+
from google.auth.transport import requests as google_requests
|
32
|
+
from google.cloud.managedkafka_v1 import Cluster, ConsumerGroup, ManagedKafkaClient, Topic, types
|
33
|
+
|
26
34
|
from airflow.exceptions import AirflowException
|
27
35
|
from airflow.providers.google.common.consts import CLIENT_INFO
|
28
36
|
from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
|
29
|
-
from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
|
30
|
-
from google.cloud.managedkafka_v1 import Cluster, ManagedKafkaClient, Topic, types
|
31
37
|
|
32
38
|
if TYPE_CHECKING:
|
33
39
|
from google.api_core.operation import Operation
|
34
40
|
from google.api_core.retry import Retry
|
35
|
-
from google.
|
41
|
+
from google.auth.credentials import Credentials
|
42
|
+
from google.cloud.managedkafka_v1.services.managed_kafka.pagers import (
|
43
|
+
ListClustersPager,
|
44
|
+
ListConsumerGroupsPager,
|
45
|
+
ListTopicsPager,
|
46
|
+
)
|
36
47
|
from google.protobuf.field_mask_pb2 import FieldMask
|
37
48
|
|
38
49
|
|
50
|
+
class ManagedKafkaTokenProvider:
|
51
|
+
"""Helper for providing authentication token for establishing connection via confluent to Apache Kafka cluster managed by Google Cloud."""
|
52
|
+
|
53
|
+
def __init__(
|
54
|
+
self,
|
55
|
+
credentials: Credentials,
|
56
|
+
):
|
57
|
+
self._credentials = credentials
|
58
|
+
self._header = json.dumps(dict(typ="JWT", alg="GOOG_OAUTH2_TOKEN"))
|
59
|
+
|
60
|
+
def _valid_credentials(self):
|
61
|
+
if not self._credentials.valid:
|
62
|
+
self._credentials.refresh(google_requests.Request())
|
63
|
+
return self._credentials
|
64
|
+
|
65
|
+
def _get_jwt(self, credentials):
|
66
|
+
return json.dumps(
|
67
|
+
dict(
|
68
|
+
exp=credentials.expiry.timestamp(),
|
69
|
+
iss="Google",
|
70
|
+
iat=datetime.datetime.now(datetime.timezone.utc).timestamp(),
|
71
|
+
scope="kafka",
|
72
|
+
sub=credentials.service_account_email,
|
73
|
+
)
|
74
|
+
)
|
75
|
+
|
76
|
+
def _b64_encode(self, source):
|
77
|
+
return base64.urlsafe_b64encode(source.encode("utf-8")).decode("utf-8").rstrip("=")
|
78
|
+
|
79
|
+
def _get_kafka_access_token(self, credentials):
|
80
|
+
return ".".join(
|
81
|
+
[
|
82
|
+
self._b64_encode(self._header),
|
83
|
+
self._b64_encode(self._get_jwt(credentials)),
|
84
|
+
self._b64_encode(credentials.token),
|
85
|
+
]
|
86
|
+
)
|
87
|
+
|
88
|
+
def confluent_token(self):
|
89
|
+
credentials = self._valid_credentials()
|
90
|
+
|
91
|
+
utc_expiry = credentials.expiry.replace(tzinfo=datetime.timezone.utc)
|
92
|
+
expiry_seconds = (utc_expiry - datetime.datetime.now(datetime.timezone.utc)).total_seconds()
|
93
|
+
|
94
|
+
return self._get_kafka_access_token(credentials), time.time() + expiry_seconds
|
95
|
+
|
96
|
+
|
39
97
|
class ManagedKafkaHook(GoogleBaseHook):
|
40
98
|
"""Hook for Managed Service for Apache Kafka APIs."""
|
41
99
|
|
@@ -62,6 +120,12 @@ class ManagedKafkaHook(GoogleBaseHook):
|
|
62
120
|
error = operation.exception(timeout=timeout)
|
63
121
|
raise AirflowException(error)
|
64
122
|
|
123
|
+
def get_confluent_token(self):
|
124
|
+
"""Get the authentication token for confluent client."""
|
125
|
+
token_provider = ManagedKafkaTokenProvider(credentials=self.get_credentials())
|
126
|
+
token = token_provider.confluent_token()
|
127
|
+
return token
|
128
|
+
|
65
129
|
@GoogleBaseHook.fallback_to_default_project_id
|
66
130
|
def create_cluster(
|
67
131
|
self,
|
@@ -480,3 +544,163 @@ class ManagedKafkaHook(GoogleBaseHook):
|
|
480
544
|
timeout=timeout,
|
481
545
|
metadata=metadata,
|
482
546
|
)
|
547
|
+
|
548
|
+
@GoogleBaseHook.fallback_to_default_project_id
|
549
|
+
def list_consumer_groups(
|
550
|
+
self,
|
551
|
+
project_id: str,
|
552
|
+
location: str,
|
553
|
+
cluster_id: str,
|
554
|
+
page_size: int | None = None,
|
555
|
+
page_token: str | None = None,
|
556
|
+
retry: Retry | _MethodDefault = DEFAULT,
|
557
|
+
timeout: float | None = None,
|
558
|
+
metadata: Sequence[tuple[str, str]] = (),
|
559
|
+
) -> ListConsumerGroupsPager:
|
560
|
+
"""
|
561
|
+
List the consumer groups in a given cluster.
|
562
|
+
|
563
|
+
:param project_id: Required. The ID of the Google Cloud project that the service belongs to.
|
564
|
+
:param location: Required. The ID of the Google Cloud region that the service belongs to.
|
565
|
+
:param cluster_id: Required. The ID of the cluster whose consumer groups are to be listed.
|
566
|
+
:param page_size: Optional. The maximum number of consumer groups to return. The service may return
|
567
|
+
fewer than this value. If unset or zero, all consumer groups for the parent is returned.
|
568
|
+
:param page_token: Optional. A page token, received from a previous ``ListConsumerGroups`` call.
|
569
|
+
Provide this to retrieve the subsequent page. When paginating, all other parameters provided to
|
570
|
+
``ListConsumerGroups`` must match the call that provided the page token.
|
571
|
+
:param retry: Designation of what errors, if any, should be retried.
|
572
|
+
:param timeout: The timeout for this request.
|
573
|
+
:param metadata: Strings which should be sent along with the request as metadata.
|
574
|
+
"""
|
575
|
+
client = self.get_managed_kafka_client()
|
576
|
+
parent = client.cluster_path(project_id, location, cluster_id)
|
577
|
+
|
578
|
+
result = client.list_consumer_groups(
|
579
|
+
request={
|
580
|
+
"parent": parent,
|
581
|
+
"page_size": page_size,
|
582
|
+
"page_token": page_token,
|
583
|
+
},
|
584
|
+
retry=retry,
|
585
|
+
timeout=timeout,
|
586
|
+
metadata=metadata,
|
587
|
+
)
|
588
|
+
return result
|
589
|
+
|
590
|
+
@GoogleBaseHook.fallback_to_default_project_id
|
591
|
+
def get_consumer_group(
|
592
|
+
self,
|
593
|
+
project_id: str,
|
594
|
+
location: str,
|
595
|
+
cluster_id: str,
|
596
|
+
consumer_group_id: str,
|
597
|
+
retry: Retry | _MethodDefault = DEFAULT,
|
598
|
+
timeout: float | None = None,
|
599
|
+
metadata: Sequence[tuple[str, str]] = (),
|
600
|
+
) -> types.ConsumerGroup:
|
601
|
+
"""
|
602
|
+
Return the properties of a single consumer group.
|
603
|
+
|
604
|
+
:param project_id: Required. The ID of the Google Cloud project that the service belongs to.
|
605
|
+
:param location: Required. The ID of the Google Cloud region that the service belongs to.
|
606
|
+
:param cluster_id: Required. The ID of the cluster whose consumer group is to be returned.
|
607
|
+
:param consumer_group_id: Required. The ID of the consumer group whose configuration to return.
|
608
|
+
:param retry: Designation of what errors, if any, should be retried.
|
609
|
+
:param timeout: The timeout for this request.
|
610
|
+
:param metadata: Strings which should be sent along with the request as metadata.
|
611
|
+
"""
|
612
|
+
client = self.get_managed_kafka_client()
|
613
|
+
name = client.consumer_group_path(project_id, location, cluster_id, consumer_group_id)
|
614
|
+
|
615
|
+
result = client.get_consumer_group(
|
616
|
+
request={
|
617
|
+
"name": name,
|
618
|
+
},
|
619
|
+
retry=retry,
|
620
|
+
timeout=timeout,
|
621
|
+
metadata=metadata,
|
622
|
+
)
|
623
|
+
return result
|
624
|
+
|
625
|
+
@GoogleBaseHook.fallback_to_default_project_id
|
626
|
+
def update_consumer_group(
|
627
|
+
self,
|
628
|
+
project_id: str,
|
629
|
+
location: str,
|
630
|
+
cluster_id: str,
|
631
|
+
consumer_group_id: str,
|
632
|
+
consumer_group: types.ConsumerGroup | dict,
|
633
|
+
update_mask: FieldMask | dict,
|
634
|
+
retry: Retry | _MethodDefault = DEFAULT,
|
635
|
+
timeout: float | None = None,
|
636
|
+
metadata: Sequence[tuple[str, str]] = (),
|
637
|
+
) -> types.ConsumerGroup:
|
638
|
+
"""
|
639
|
+
Update the properties of a single consumer group.
|
640
|
+
|
641
|
+
:param project_id: Required. The ID of the Google Cloud project that the service belongs to.
|
642
|
+
:param location: Required. The ID of the Google Cloud region that the service belongs to.
|
643
|
+
:param cluster_id: Required. The ID of the cluster whose topic is to be updated.
|
644
|
+
:param consumer_group_id: Required. The ID of the consumer group whose configuration to update.
|
645
|
+
:param consumer_group: Required. The consumer_group to update. Its ``name`` field must be populated.
|
646
|
+
:param update_mask: Required. Field mask is used to specify the fields to be overwritten in the
|
647
|
+
ConsumerGroup resource by the update. The fields specified in the update_mask are relative to the
|
648
|
+
resource, not the full request. A field will be overwritten if it is in the mask.
|
649
|
+
:param retry: Designation of what errors, if any, should be retried.
|
650
|
+
:param timeout: The timeout for this request.
|
651
|
+
:param metadata: Strings which should be sent along with the request as metadata.
|
652
|
+
"""
|
653
|
+
client = self.get_managed_kafka_client()
|
654
|
+
_consumer_group = (
|
655
|
+
deepcopy(consumer_group)
|
656
|
+
if isinstance(consumer_group, dict)
|
657
|
+
else ConsumerGroup.to_dict(consumer_group)
|
658
|
+
)
|
659
|
+
_consumer_group["name"] = client.consumer_group_path(
|
660
|
+
project_id, location, cluster_id, consumer_group_id
|
661
|
+
)
|
662
|
+
|
663
|
+
result = client.update_consumer_group(
|
664
|
+
request={
|
665
|
+
"update_mask": update_mask,
|
666
|
+
"consumer_group": _consumer_group,
|
667
|
+
},
|
668
|
+
retry=retry,
|
669
|
+
timeout=timeout,
|
670
|
+
metadata=metadata,
|
671
|
+
)
|
672
|
+
return result
|
673
|
+
|
674
|
+
@GoogleBaseHook.fallback_to_default_project_id
|
675
|
+
def delete_consumer_group(
|
676
|
+
self,
|
677
|
+
project_id: str,
|
678
|
+
location: str,
|
679
|
+
cluster_id: str,
|
680
|
+
consumer_group_id: str,
|
681
|
+
retry: Retry | _MethodDefault = DEFAULT,
|
682
|
+
timeout: float | None = None,
|
683
|
+
metadata: Sequence[tuple[str, str]] = (),
|
684
|
+
) -> None:
|
685
|
+
"""
|
686
|
+
Delete a single consumer group.
|
687
|
+
|
688
|
+
:param project_id: Required. The ID of the Google Cloud project that the service belongs to.
|
689
|
+
:param location: Required. The ID of the Google Cloud region that the service belongs to.
|
690
|
+
:param cluster_id: Required. The ID of the cluster whose consumer group is to be deleted.
|
691
|
+
:param consumer_group_id: Required. The ID of the consumer group to delete.
|
692
|
+
:param retry: Designation of what errors, if any, should be retried.
|
693
|
+
:param timeout: The timeout for this request.
|
694
|
+
:param metadata: Strings which should be sent along with the request as metadata.
|
695
|
+
"""
|
696
|
+
client = self.get_managed_kafka_client()
|
697
|
+
name = client.consumer_group_path(project_id, location, cluster_id, consumer_group_id)
|
698
|
+
|
699
|
+
client.delete_consumer_group(
|
700
|
+
request={
|
701
|
+
"name": name,
|
702
|
+
},
|
703
|
+
retry=retry,
|
704
|
+
timeout=timeout,
|
705
|
+
metadata=metadata,
|
706
|
+
)
|
@@ -22,8 +22,6 @@ from __future__ import annotations
|
|
22
22
|
from collections.abc import Sequence
|
23
23
|
from typing import TYPE_CHECKING
|
24
24
|
|
25
|
-
from airflow.providers.google.common.consts import CLIENT_INFO
|
26
|
-
from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
|
27
25
|
from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
|
28
26
|
from google.cloud.language_v1 import EncodingType, LanguageServiceClient
|
29
27
|
from google.cloud.language_v1.types import (
|
@@ -37,6 +35,9 @@ from google.cloud.language_v1.types import (
|
|
37
35
|
Document,
|
38
36
|
)
|
39
37
|
|
38
|
+
from airflow.providers.google.common.consts import CLIENT_INFO
|
39
|
+
from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
|
40
|
+
|
40
41
|
if TYPE_CHECKING:
|
41
42
|
from google.api_core.retry import Retry
|
42
43
|
|