apache-airflow-providers-google 10.7.0__py3-none-any.whl → 10.8.0__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.
Files changed (174) hide show
  1. airflow/providers/google/__init__.py +1 -1
  2. airflow/providers/google/ads/hooks/ads.py +10 -6
  3. airflow/providers/google/cloud/_internal_client/secret_manager_client.py +4 -1
  4. airflow/providers/google/cloud/example_dags/example_cloud_sql_query.py +31 -34
  5. airflow/providers/google/cloud/hooks/automl.py +11 -9
  6. airflow/providers/google/cloud/hooks/bigquery.py +30 -36
  7. airflow/providers/google/cloud/hooks/bigquery_dts.py +5 -3
  8. airflow/providers/google/cloud/hooks/bigtable.py +11 -8
  9. airflow/providers/google/cloud/hooks/cloud_batch.py +5 -3
  10. airflow/providers/google/cloud/hooks/cloud_build.py +6 -4
  11. airflow/providers/google/cloud/hooks/cloud_composer.py +14 -10
  12. airflow/providers/google/cloud/hooks/cloud_memorystore.py +5 -3
  13. airflow/providers/google/cloud/hooks/cloud_run.py +5 -3
  14. airflow/providers/google/cloud/hooks/cloud_sql.py +11 -14
  15. airflow/providers/google/cloud/hooks/cloud_storage_transfer_service.py +8 -6
  16. airflow/providers/google/cloud/hooks/compute.py +5 -3
  17. airflow/providers/google/cloud/hooks/compute_ssh.py +1 -1
  18. airflow/providers/google/cloud/hooks/datacatalog.py +5 -3
  19. airflow/providers/google/cloud/hooks/dataflow.py +8 -11
  20. airflow/providers/google/cloud/hooks/dataform.py +4 -2
  21. airflow/providers/google/cloud/hooks/datafusion.py +24 -6
  22. airflow/providers/google/cloud/hooks/dataplex.py +75 -6
  23. airflow/providers/google/cloud/hooks/dataproc.py +9 -7
  24. airflow/providers/google/cloud/hooks/dataproc_metastore.py +8 -6
  25. airflow/providers/google/cloud/hooks/dlp.py +139 -137
  26. airflow/providers/google/cloud/hooks/gcs.py +15 -20
  27. airflow/providers/google/cloud/hooks/kms.py +4 -2
  28. airflow/providers/google/cloud/hooks/kubernetes_engine.py +34 -34
  29. airflow/providers/google/cloud/hooks/looker.py +4 -1
  30. airflow/providers/google/cloud/hooks/mlengine.py +8 -6
  31. airflow/providers/google/cloud/hooks/natural_language.py +4 -2
  32. airflow/providers/google/cloud/hooks/os_login.py +9 -7
  33. airflow/providers/google/cloud/hooks/pubsub.py +13 -11
  34. airflow/providers/google/cloud/hooks/spanner.py +7 -5
  35. airflow/providers/google/cloud/hooks/speech_to_text.py +4 -2
  36. airflow/providers/google/cloud/hooks/stackdriver.py +6 -5
  37. airflow/providers/google/cloud/hooks/tasks.py +5 -3
  38. airflow/providers/google/cloud/hooks/text_to_speech.py +4 -2
  39. airflow/providers/google/cloud/hooks/vertex_ai/auto_ml.py +7 -5
  40. airflow/providers/google/cloud/hooks/vertex_ai/batch_prediction_job.py +6 -4
  41. airflow/providers/google/cloud/hooks/vertex_ai/custom_job.py +11 -9
  42. airflow/providers/google/cloud/hooks/vertex_ai/dataset.py +12 -10
  43. airflow/providers/google/cloud/hooks/vertex_ai/endpoint_service.py +8 -6
  44. airflow/providers/google/cloud/hooks/vertex_ai/hyperparameter_tuning_job.py +6 -4
  45. airflow/providers/google/cloud/hooks/vertex_ai/model_service.py +7 -5
  46. airflow/providers/google/cloud/hooks/video_intelligence.py +5 -3
  47. airflow/providers/google/cloud/hooks/vision.py +5 -3
  48. airflow/providers/google/cloud/hooks/workflows.py +8 -6
  49. airflow/providers/google/cloud/links/bigquery.py +1 -1
  50. airflow/providers/google/cloud/links/bigquery_dts.py +1 -1
  51. airflow/providers/google/cloud/links/cloud_functions.py +1 -1
  52. airflow/providers/google/cloud/links/cloud_memorystore.py +1 -1
  53. airflow/providers/google/cloud/links/cloud_sql.py +1 -1
  54. airflow/providers/google/cloud/links/cloud_tasks.py +1 -1
  55. airflow/providers/google/cloud/links/compute.py +1 -1
  56. airflow/providers/google/cloud/links/datacatalog.py +1 -1
  57. airflow/providers/google/cloud/links/dataflow.py +1 -1
  58. airflow/providers/google/cloud/links/dataform.py +1 -1
  59. airflow/providers/google/cloud/links/pubsub.py +1 -1
  60. airflow/providers/google/cloud/links/spanner.py +1 -1
  61. airflow/providers/google/cloud/links/stackdriver.py +1 -1
  62. airflow/providers/google/cloud/links/workflows.py +2 -2
  63. airflow/providers/google/cloud/log/gcs_task_handler.py +5 -7
  64. airflow/providers/google/cloud/log/stackdriver_task_handler.py +8 -4
  65. airflow/providers/google/cloud/operators/automl.py +2 -1
  66. airflow/providers/google/cloud/operators/bigquery.py +6 -2
  67. airflow/providers/google/cloud/operators/bigquery_dts.py +2 -1
  68. airflow/providers/google/cloud/operators/bigtable.py +5 -3
  69. airflow/providers/google/cloud/operators/cloud_batch.py +6 -3
  70. airflow/providers/google/cloud/operators/cloud_build.py +2 -1
  71. airflow/providers/google/cloud/operators/cloud_composer.py +3 -2
  72. airflow/providers/google/cloud/operators/cloud_memorystore.py +3 -2
  73. airflow/providers/google/cloud/operators/cloud_run.py +3 -2
  74. airflow/providers/google/cloud/operators/cloud_sql.py +157 -152
  75. airflow/providers/google/cloud/operators/compute.py +59 -61
  76. airflow/providers/google/cloud/operators/datacatalog.py +3 -2
  77. airflow/providers/google/cloud/operators/dataflow.py +3 -1
  78. airflow/providers/google/cloud/operators/dataform.py +2 -1
  79. airflow/providers/google/cloud/operators/datafusion.py +1 -1
  80. airflow/providers/google/cloud/operators/dataplex.py +110 -8
  81. airflow/providers/google/cloud/operators/dataproc.py +39 -18
  82. airflow/providers/google/cloud/operators/dataproc_metastore.py +2 -1
  83. airflow/providers/google/cloud/operators/dlp.py +3 -2
  84. airflow/providers/google/cloud/operators/functions.py +46 -46
  85. airflow/providers/google/cloud/operators/gcs.py +4 -6
  86. airflow/providers/google/cloud/operators/kubernetes_engine.py +2 -1
  87. airflow/providers/google/cloud/operators/natural_language.py +3 -2
  88. airflow/providers/google/cloud/operators/pubsub.py +2 -1
  89. airflow/providers/google/cloud/operators/speech_to_text.py +3 -2
  90. airflow/providers/google/cloud/operators/stackdriver.py +2 -1
  91. airflow/providers/google/cloud/operators/tasks.py +3 -2
  92. airflow/providers/google/cloud/operators/text_to_speech.py +3 -2
  93. airflow/providers/google/cloud/operators/translate_speech.py +2 -1
  94. airflow/providers/google/cloud/operators/vertex_ai/auto_ml.py +2 -1
  95. airflow/providers/google/cloud/operators/vertex_ai/batch_prediction_job.py +3 -2
  96. airflow/providers/google/cloud/operators/vertex_ai/custom_job.py +2 -1
  97. airflow/providers/google/cloud/operators/vertex_ai/dataset.py +3 -2
  98. airflow/providers/google/cloud/operators/vertex_ai/endpoint_service.py +4 -4
  99. airflow/providers/google/cloud/operators/vertex_ai/hyperparameter_tuning_job.py +3 -2
  100. airflow/providers/google/cloud/operators/vertex_ai/model_service.py +2 -1
  101. airflow/providers/google/cloud/operators/video_intelligence.py +2 -1
  102. airflow/providers/google/cloud/operators/vision.py +3 -2
  103. airflow/providers/google/cloud/operators/workflows.py +7 -5
  104. airflow/providers/google/cloud/secrets/secret_manager.py +2 -2
  105. airflow/providers/google/cloud/sensors/bigquery_dts.py +2 -1
  106. airflow/providers/google/cloud/sensors/dataplex.py +2 -1
  107. airflow/providers/google/cloud/sensors/dataproc_metastore.py +2 -2
  108. airflow/providers/google/cloud/sensors/gcs.py +2 -1
  109. airflow/providers/google/cloud/sensors/workflows.py +2 -1
  110. airflow/providers/google/cloud/transfers/azure_fileshare_to_gcs.py +24 -10
  111. airflow/providers/google/cloud/transfers/bigquery_to_gcs.py +2 -1
  112. airflow/providers/google/cloud/transfers/bigquery_to_mssql.py +2 -1
  113. airflow/providers/google/cloud/transfers/bigquery_to_mysql.py +1 -4
  114. airflow/providers/google/cloud/transfers/bigquery_to_postgres.py +1 -4
  115. airflow/providers/google/cloud/transfers/bigquery_to_sql.py +1 -1
  116. airflow/providers/google/cloud/transfers/calendar_to_gcs.py +4 -2
  117. airflow/providers/google/cloud/transfers/cassandra_to_gcs.py +1 -3
  118. airflow/providers/google/cloud/transfers/facebook_ads_to_gcs.py +2 -2
  119. airflow/providers/google/cloud/transfers/gcs_to_bigquery.py +2 -1
  120. airflow/providers/google/cloud/transfers/presto_to_gcs.py +5 -4
  121. airflow/providers/google/cloud/transfers/sql_to_gcs.py +1 -1
  122. airflow/providers/google/cloud/transfers/trino_to_gcs.py +5 -4
  123. airflow/providers/google/cloud/triggers/bigquery.py +30 -36
  124. airflow/providers/google/cloud/triggers/bigquery_dts.py +9 -10
  125. airflow/providers/google/cloud/triggers/cloud_batch.py +6 -8
  126. airflow/providers/google/cloud/triggers/cloud_build.py +5 -6
  127. airflow/providers/google/cloud/triggers/cloud_run.py +4 -3
  128. airflow/providers/google/cloud/triggers/cloud_sql.py +10 -10
  129. airflow/providers/google/cloud/triggers/cloud_storage_transfer_service.py +1 -1
  130. airflow/providers/google/cloud/triggers/dataflow.py +5 -6
  131. airflow/providers/google/cloud/triggers/datafusion.py +5 -6
  132. airflow/providers/google/cloud/triggers/dataplex.py +110 -0
  133. airflow/providers/google/cloud/triggers/dataproc.py +18 -20
  134. airflow/providers/google/cloud/triggers/kubernetes_engine.py +14 -13
  135. airflow/providers/google/cloud/triggers/mlengine.py +5 -5
  136. airflow/providers/google/cloud/triggers/pubsub.py +2 -2
  137. airflow/providers/google/cloud/utils/bigquery_get_data.py +6 -3
  138. airflow/providers/google/cloud/utils/credentials_provider.py +1 -1
  139. airflow/providers/google/cloud/utils/field_validator.py +13 -13
  140. airflow/providers/google/cloud/utils/mlengine_operator_utils.py +5 -3
  141. airflow/providers/google/cloud/utils/mlengine_prediction_summary.py +1 -1
  142. airflow/providers/google/common/hooks/base_google.py +10 -2
  143. airflow/providers/google/common/links/storage.py +1 -1
  144. airflow/providers/google/common/utils/id_token_credentials.py +4 -1
  145. airflow/providers/google/get_provider_info.py +5 -0
  146. airflow/providers/google/marketing_platform/hooks/campaign_manager.py +4 -2
  147. airflow/providers/google/marketing_platform/sensors/display_video.py +6 -3
  148. airflow/providers/google/suite/hooks/calendar.py +4 -2
  149. {apache_airflow_providers_google-10.7.0.dist-info → apache_airflow_providers_google-10.8.0.dist-info}/METADATA +6 -6
  150. {apache_airflow_providers_google-10.7.0.dist-info → apache_airflow_providers_google-10.8.0.dist-info}/RECORD +155 -173
  151. {apache_airflow_providers_google-10.7.0.dist-info → apache_airflow_providers_google-10.8.0.dist-info}/WHEEL +1 -1
  152. airflow/providers/google/ads/_vendor/__init__.py +0 -16
  153. airflow/providers/google/ads/_vendor/googleads/__init__.py +0 -16
  154. airflow/providers/google/ads/_vendor/googleads/interceptors/__init__.py +0 -16
  155. airflow/providers/google/ads/_vendor/googleads/v12/__init__.py +0 -16
  156. airflow/providers/google/ads/_vendor/googleads/v12/common/__init__.py +0 -16
  157. airflow/providers/google/ads/_vendor/googleads/v12/common/types/__init__.py +0 -16
  158. airflow/providers/google/ads/_vendor/googleads/v12/enums/__init__.py +0 -16
  159. airflow/providers/google/ads/_vendor/googleads/v12/enums/types/__init__.py +0 -16
  160. airflow/providers/google/ads/_vendor/googleads/v12/errors/__init__.py +0 -16
  161. airflow/providers/google/ads/_vendor/googleads/v12/errors/types/__init__.py +0 -16
  162. airflow/providers/google/ads/_vendor/googleads/v12/resources/__init__.py +0 -16
  163. airflow/providers/google/ads/_vendor/googleads/v12/resources/types/__init__.py +0 -16
  164. airflow/providers/google/ads/_vendor/googleads/v12/services/__init__.py +0 -16
  165. airflow/providers/google/ads/_vendor/googleads/v12/services/services/__init__.py +0 -16
  166. airflow/providers/google/ads/_vendor/googleads/v12/services/services/customer_service/__init__.py +0 -16
  167. airflow/providers/google/ads/_vendor/googleads/v12/services/services/customer_service/transports/__init__.py +0 -16
  168. airflow/providers/google/ads/_vendor/googleads/v12/services/services/google_ads_service/__init__.py +0 -16
  169. airflow/providers/google/ads/_vendor/googleads/v12/services/services/google_ads_service/transports/__init__.py +0 -16
  170. airflow/providers/google/ads/_vendor/googleads/v12/services/types/__init__.py +0 -16
  171. {apache_airflow_providers_google-10.7.0.dist-info → apache_airflow_providers_google-10.8.0.dist-info}/LICENSE +0 -0
  172. {apache_airflow_providers_google-10.7.0.dist-info → apache_airflow_providers_google-10.8.0.dist-info}/NOTICE +0 -0
  173. {apache_airflow_providers_google-10.7.0.dist-info → apache_airflow_providers_google-10.8.0.dist-info}/entry_points.txt +0 -0
  174. {apache_airflow_providers_google-10.7.0.dist-info → apache_airflow_providers_google-10.8.0.dist-info}/top_level.txt +0 -0
@@ -30,13 +30,11 @@ import json
30
30
  import time
31
31
  import warnings
32
32
  from functools import cached_property
33
- from typing import Sequence
33
+ from typing import TYPE_CHECKING, Sequence
34
34
 
35
- import google.auth.credentials
36
35
  from gcloud.aio.auth import Token
37
36
  from google.api_core.exceptions import NotFound
38
37
  from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
39
- from google.api_core.retry import Retry
40
38
  from google.auth.transport import requests as google_requests
41
39
 
42
40
  # not sure why but mypy complains on missing `container_v1` but it is clearly there and is importable
@@ -45,7 +43,6 @@ from google.cloud.container_v1 import ClusterManagerAsyncClient, ClusterManagerC
45
43
  from google.cloud.container_v1.types import Cluster, Operation
46
44
  from kubernetes import client
47
45
  from kubernetes_asyncio import client as async_client
48
- from kubernetes_asyncio.client.models import V1Pod
49
46
  from kubernetes_asyncio.config.kube_config import FileOrData
50
47
  from urllib3.exceptions import HTTPError
51
48
 
@@ -59,6 +56,11 @@ from airflow.providers.google.common.hooks.base_google import (
59
56
  GoogleBaseHook,
60
57
  )
61
58
 
59
+ if TYPE_CHECKING:
60
+ import google.auth.credentials
61
+ from google.api_core.retry import Retry
62
+ from kubernetes_asyncio.client.models import V1Pod
63
+
62
64
  OPERATIONAL_POLL_INTERVAL = 15
63
65
 
64
66
 
@@ -493,19 +495,18 @@ class GKEPodAsyncHook(GoogleBaseAsyncHook):
493
495
  :param name: Name of the pod.
494
496
  :param namespace: Name of the pod's namespace.
495
497
  """
496
- async with Token(scopes=self.scopes) as token:
497
- async with self.get_conn(token) as connection:
498
- try:
499
- v1_api = async_client.CoreV1Api(connection)
500
- await v1_api.delete_namespaced_pod(
501
- name=name,
502
- namespace=namespace,
503
- body=client.V1DeleteOptions(),
504
- )
505
- except async_client.ApiException as e:
506
- # If the pod is already deleted
507
- if e.status != 404:
508
- raise
498
+ async with Token(scopes=self.scopes) as token, self.get_conn(token) as connection:
499
+ try:
500
+ v1_api = async_client.CoreV1Api(connection)
501
+ await v1_api.delete_namespaced_pod(
502
+ name=name,
503
+ namespace=namespace,
504
+ body=client.V1DeleteOptions(),
505
+ )
506
+ except async_client.ApiException as e:
507
+ # If the pod is already deleted
508
+ if e.status != 404:
509
+ raise
509
510
 
510
511
  async def read_logs(self, name: str, namespace: str):
511
512
  """Read logs inside the pod while starting containers inside.
@@ -518,20 +519,19 @@ class GKEPodAsyncHook(GoogleBaseAsyncHook):
518
519
  :param name: Name of the pod.
519
520
  :param namespace: Name of the pod's namespace.
520
521
  """
521
- async with Token(scopes=self.scopes) as token:
522
- async with self.get_conn(token) as connection:
523
- try:
524
- v1_api = async_client.CoreV1Api(connection)
525
- logs = await v1_api.read_namespaced_pod_log(
526
- name=name,
527
- namespace=namespace,
528
- follow=False,
529
- timestamps=True,
530
- )
531
- logs = logs.splitlines()
532
- for line in logs:
533
- self.log.info("Container logs from %s", line)
534
- return logs
535
- except HTTPError:
536
- self.log.exception("There was an error reading the kubernetes API.")
537
- raise
522
+ async with Token(scopes=self.scopes) as token, self.get_conn(token) as connection:
523
+ try:
524
+ v1_api = async_client.CoreV1Api(connection)
525
+ logs = await v1_api.read_namespaced_pod_log(
526
+ name=name,
527
+ namespace=namespace,
528
+ follow=False,
529
+ timestamps=True,
530
+ )
531
+ logs = logs.splitlines()
532
+ for line in logs:
533
+ self.log.info("Container logs from %s", line)
534
+ return logs
535
+ except HTTPError:
536
+ self.log.exception("There was an error reading the kubernetes API.")
537
+ raise
@@ -21,6 +21,7 @@ from __future__ import annotations
21
21
  import json
22
22
  import time
23
23
  from enum import Enum
24
+ from typing import TYPE_CHECKING
24
25
 
25
26
  from looker_sdk.rtl import api_settings, auth_session, requests_transport, serialize
26
27
  from looker_sdk.sdk.api40 import methods as methods40
@@ -28,9 +29,11 @@ from packaging.version import parse as parse_version
28
29
 
29
30
  from airflow.exceptions import AirflowException
30
31
  from airflow.hooks.base import BaseHook
31
- from airflow.models.connection import Connection
32
32
  from airflow.version import version
33
33
 
34
+ if TYPE_CHECKING:
35
+ from airflow.models.connection import Connection
36
+
34
37
 
35
38
  class LookerHook(BaseHook):
36
39
  """Hook for Looker APIs."""
@@ -18,22 +18,25 @@
18
18
  """This module contains a Google ML Engine Hook."""
19
19
  from __future__ import annotations
20
20
 
21
+ import contextlib
21
22
  import logging
22
23
  import random
23
24
  import time
24
- from typing import Callable
25
+ from typing import TYPE_CHECKING, Callable
25
26
 
26
27
  from aiohttp import ClientSession
27
28
  from gcloud.aio.auth import AioSession, Token
28
29
  from googleapiclient.discovery import Resource, build
29
30
  from googleapiclient.errors import HttpError
30
- from httplib2 import Response
31
- from requests import Session
32
31
 
33
32
  from airflow.exceptions import AirflowException
34
33
  from airflow.providers.google.common.hooks.base_google import GoogleBaseAsyncHook, GoogleBaseHook
35
34
  from airflow.version import version as airflow_version
36
35
 
36
+ if TYPE_CHECKING:
37
+ from httplib2 import Response
38
+ from requests import Session
39
+
37
40
  log = logging.getLogger(__name__)
38
41
 
39
42
  _AIRFLOW_VERSION = "v" + airflow_version.replace(".", "-").replace("+", "-")
@@ -559,10 +562,9 @@ class MLEngineAsyncHook(GoogleBaseAsyncHook):
559
562
  headers = {
560
563
  "Authorization": f"Bearer {await token.get()}",
561
564
  }
562
- try:
565
+ with contextlib.suppress(AirflowException):
566
+ # suppress AirflowException because we don't want to raise exception
563
567
  job = await session_aio.get(url=url, headers=headers)
564
- except AirflowException:
565
- pass # Because the job may not be visible in system yet
566
568
 
567
569
  return job
568
570
 
@@ -18,10 +18,9 @@
18
18
  """This module contains a Google Cloud Natural Language Hook."""
19
19
  from __future__ import annotations
20
20
 
21
- from typing import Sequence
21
+ from typing import TYPE_CHECKING, Sequence
22
22
 
23
23
  from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
24
- from google.api_core.retry import Retry
25
24
  from google.cloud.language_v1 import EncodingType, LanguageServiceClient
26
25
  from google.cloud.language_v1.types import (
27
26
  AnalyzeEntitiesResponse,
@@ -37,6 +36,9 @@ from google.cloud.language_v1.types import (
37
36
  from airflow.providers.google.common.consts import CLIENT_INFO
38
37
  from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
39
38
 
39
+ if TYPE_CHECKING:
40
+ from google.api_core.retry import Retry
41
+
40
42
 
41
43
  class CloudNaturalLanguageHook(GoogleBaseHook):
42
44
  """
@@ -22,15 +22,17 @@
22
22
  """
23
23
  from __future__ import annotations
24
24
 
25
- from typing import Sequence
25
+ from typing import TYPE_CHECKING, Sequence
26
26
 
27
27
  from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
28
- from google.api_core.retry import Retry
29
28
  from google.cloud.oslogin_v1 import ImportSshPublicKeyResponse, OsLoginServiceClient
30
29
 
31
30
  from airflow.providers.google.common.consts import CLIENT_INFO
32
31
  from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID, GoogleBaseHook
33
32
 
33
+ if TYPE_CHECKING:
34
+ from google.api_core.retry import Retry
35
+
34
36
 
35
37
  class OSLoginHook(GoogleBaseHook):
36
38
  """
@@ -92,11 +94,11 @@ class OSLoginHook(GoogleBaseHook):
92
94
  """
93
95
  conn = self.get_conn()
94
96
  return conn.import_ssh_public_key(
95
- request=dict(
96
- parent=f"users/{user}",
97
- ssh_public_key=ssh_public_key,
98
- project_id=project_id,
99
- ),
97
+ request={
98
+ "parent": f"users/{user}",
99
+ "ssh_public_key": ssh_public_key,
100
+ "project_id": project_id,
101
+ },
100
102
  retry=retry,
101
103
  timeout=timeout,
102
104
  metadata=metadata,
@@ -28,23 +28,13 @@ from __future__ import annotations
28
28
  import warnings
29
29
  from base64 import b64decode
30
30
  from functools import cached_property
31
- from typing import Any, Sequence
31
+ from typing import TYPE_CHECKING, Any, Sequence
32
32
  from uuid import uuid4
33
33
 
34
34
  from google.api_core.exceptions import AlreadyExists, GoogleAPICallError
35
35
  from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
36
- from google.api_core.retry import Retry
37
36
  from google.cloud.exceptions import NotFound
38
37
  from google.cloud.pubsub_v1 import PublisherClient, SubscriberClient
39
- from google.cloud.pubsub_v1.types import (
40
- DeadLetterPolicy,
41
- Duration,
42
- ExpirationPolicy,
43
- MessageStoragePolicy,
44
- PushConfig,
45
- ReceivedMessage,
46
- RetryPolicy,
47
- )
48
38
  from google.pubsub_v1.services.subscriber.async_client import SubscriberAsyncClient
49
39
  from googleapiclient.errors import HttpError
50
40
 
@@ -57,6 +47,18 @@ from airflow.providers.google.common.hooks.base_google import (
57
47
  )
58
48
  from airflow.version import version
59
49
 
50
+ if TYPE_CHECKING:
51
+ from google.api_core.retry import Retry
52
+ from google.cloud.pubsub_v1.types import (
53
+ DeadLetterPolicy,
54
+ Duration,
55
+ ExpirationPolicy,
56
+ MessageStoragePolicy,
57
+ PushConfig,
58
+ ReceivedMessage,
59
+ RetryPolicy,
60
+ )
61
+
60
62
 
61
63
  class PubSubException(Exception):
62
64
  """Alias for Exception."""
@@ -18,14 +18,10 @@
18
18
  """This module contains a Google Cloud Spanner Hook."""
19
19
  from __future__ import annotations
20
20
 
21
- from typing import Callable, NamedTuple, Sequence
21
+ from typing import TYPE_CHECKING, Callable, NamedTuple, Sequence
22
22
 
23
23
  from google.api_core.exceptions import AlreadyExists, GoogleAPICallError
24
24
  from google.cloud.spanner_v1.client import Client
25
- from google.cloud.spanner_v1.database import Database
26
- from google.cloud.spanner_v1.instance import Instance
27
- from google.cloud.spanner_v1.transaction import Transaction
28
- from google.longrunning.operations_grpc_pb2 import Operation
29
25
  from sqlalchemy import create_engine
30
26
 
31
27
  from airflow.exceptions import AirflowException
@@ -33,6 +29,12 @@ from airflow.providers.common.sql.hooks.sql import DbApiHook
33
29
  from airflow.providers.google.common.consts import CLIENT_INFO
34
30
  from airflow.providers.google.common.hooks.base_google import GoogleBaseHook, get_field
35
31
 
32
+ if TYPE_CHECKING:
33
+ from google.cloud.spanner_v1.database import Database
34
+ from google.cloud.spanner_v1.instance import Instance
35
+ from google.cloud.spanner_v1.transaction import Transaction
36
+ from google.longrunning.operations_grpc_pb2 import Operation
37
+
36
38
 
37
39
  class SpannerConnectionParams(NamedTuple):
38
40
  """Information about Google Spanner connection parameters."""
@@ -18,16 +18,18 @@
18
18
  """This module contains a Google Cloud Speech Hook."""
19
19
  from __future__ import annotations
20
20
 
21
- from typing import Sequence
21
+ from typing import TYPE_CHECKING, Sequence
22
22
 
23
23
  from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
24
- from google.api_core.retry import Retry
25
24
  from google.cloud.speech_v1 import SpeechClient
26
25
  from google.cloud.speech_v1.types import RecognitionAudio, RecognitionConfig
27
26
 
28
27
  from airflow.providers.google.common.consts import CLIENT_INFO
29
28
  from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
30
29
 
30
+ if TYPE_CHECKING:
31
+ from google.api_core.retry import Retry
32
+
31
33
 
32
34
  class CloudSpeechToTextHook(GoogleBaseHook):
33
35
  """
@@ -18,12 +18,12 @@
18
18
  """This module contains Google Cloud Stackdriver operators."""
19
19
  from __future__ import annotations
20
20
 
21
+ import contextlib
21
22
  import json
22
- from typing import Any, Sequence
23
+ from typing import TYPE_CHECKING, Any, Sequence
23
24
 
24
25
  from google.api_core.exceptions import InvalidArgument
25
26
  from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
26
- from google.api_core.retry import Retry
27
27
  from google.cloud import monitoring_v3
28
28
  from google.cloud.monitoring_v3 import AlertPolicy, NotificationChannel
29
29
  from google.protobuf.field_mask_pb2 import FieldMask
@@ -32,6 +32,9 @@ from googleapiclient.errors import HttpError
32
32
  from airflow.exceptions import AirflowException
33
33
  from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID, GoogleBaseHook
34
34
 
35
+ if TYPE_CHECKING:
36
+ from google.api_core.retry import Retry
37
+
35
38
 
36
39
  class StackdriverHook(GoogleBaseHook):
37
40
  """Stackdriver Hook for connecting with Google Cloud Stackdriver."""
@@ -293,15 +296,13 @@ class StackdriverHook(GoogleBaseHook):
293
296
  policy.notification_channels[i] = new_channel
294
297
 
295
298
  if policy.name in existing_policies:
296
- try:
299
+ with contextlib.suppress(InvalidArgument):
297
300
  policy_client.update_alert_policy(
298
301
  request={"alert_policy": policy},
299
302
  retry=retry,
300
303
  timeout=timeout,
301
304
  metadata=metadata,
302
305
  )
303
- except InvalidArgument:
304
- pass
305
306
  else:
306
307
  policy.name = None
307
308
  for condition in policy.conditions:
@@ -19,18 +19,20 @@
19
19
 
20
20
  from __future__ import annotations
21
21
 
22
- from typing import Sequence
22
+ from typing import TYPE_CHECKING, Sequence
23
23
 
24
24
  from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
25
- from google.api_core.retry import Retry
26
25
  from google.cloud.tasks_v2 import CloudTasksClient
27
26
  from google.cloud.tasks_v2.types import Queue, Task
28
- from google.protobuf.field_mask_pb2 import FieldMask
29
27
 
30
28
  from airflow.exceptions import AirflowException
31
29
  from airflow.providers.google.common.consts import CLIENT_INFO
32
30
  from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID, GoogleBaseHook
33
31
 
32
+ if TYPE_CHECKING:
33
+ from google.api_core.retry import Retry
34
+ from google.protobuf.field_mask_pb2 import FieldMask
35
+
34
36
 
35
37
  class CloudTasksHook(GoogleBaseHook):
36
38
  """
@@ -18,10 +18,9 @@
18
18
  """This module contains a Google Cloud Text to Speech Hook."""
19
19
  from __future__ import annotations
20
20
 
21
- from typing import Sequence
21
+ from typing import TYPE_CHECKING, Sequence
22
22
 
23
23
  from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
24
- from google.api_core.retry import Retry
25
24
  from google.cloud.texttospeech_v1 import TextToSpeechClient
26
25
  from google.cloud.texttospeech_v1.types import (
27
26
  AudioConfig,
@@ -33,6 +32,9 @@ from google.cloud.texttospeech_v1.types import (
33
32
  from airflow.providers.google.common.consts import CLIENT_INFO
34
33
  from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
35
34
 
35
+ if TYPE_CHECKING:
36
+ from google.api_core.retry import Retry
37
+
36
38
 
37
39
  class CloudTextToSpeechHook(GoogleBaseHook):
38
40
  """
@@ -45,12 +45,10 @@ This module contains a Google Cloud Vertex AI hook.
45
45
  from __future__ import annotations
46
46
 
47
47
  import warnings
48
- from typing import Sequence
48
+ from typing import TYPE_CHECKING, Sequence
49
49
 
50
50
  from google.api_core.client_options import ClientOptions
51
51
  from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
52
- from google.api_core.operation import Operation
53
- from google.api_core.retry import Retry
54
52
  from google.cloud.aiplatform import (
55
53
  AutoMLForecastingTrainingJob,
56
54
  AutoMLImageTrainingJob,
@@ -61,12 +59,16 @@ from google.cloud.aiplatform import (
61
59
  models,
62
60
  )
63
61
  from google.cloud.aiplatform_v1 import JobServiceClient, PipelineServiceClient
64
- from google.cloud.aiplatform_v1.services.pipeline_service.pagers import ListTrainingPipelinesPager
65
- from google.cloud.aiplatform_v1.types import TrainingPipeline
66
62
 
67
63
  from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
68
64
  from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
69
65
 
66
+ if TYPE_CHECKING:
67
+ from google.api_core.operation import Operation
68
+ from google.api_core.retry import Retry
69
+ from google.cloud.aiplatform_v1.services.pipeline_service.pagers import ListTrainingPipelinesPager
70
+ from google.cloud.aiplatform_v1.types import TrainingPipeline
71
+
70
72
 
71
73
  class AutoMLHook(GoogleBaseHook):
72
74
  """Hook for Google Cloud Vertex AI Auto ML APIs."""
@@ -26,19 +26,21 @@
26
26
  """
27
27
  from __future__ import annotations
28
28
 
29
- from typing import Sequence
29
+ from typing import TYPE_CHECKING, Sequence
30
30
 
31
31
  from google.api_core.client_options import ClientOptions
32
32
  from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
33
- from google.api_core.operation import Operation
34
- from google.api_core.retry import Retry
35
33
  from google.cloud.aiplatform import BatchPredictionJob, Model, explain
36
34
  from google.cloud.aiplatform_v1 import JobServiceClient
37
- from google.cloud.aiplatform_v1.services.job_service.pagers import ListBatchPredictionJobsPager
38
35
 
39
36
  from airflow import AirflowException
40
37
  from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
41
38
 
39
+ if TYPE_CHECKING:
40
+ from google.api_core.operation import Operation
41
+ from google.api_core.retry import Retry
42
+ from google.cloud.aiplatform_v1.services.job_service.pagers import ListBatchPredictionJobsPager
43
+
42
44
 
43
45
  class BatchPredictionJobHook(GoogleBaseHook):
44
46
  """Hook for Google Cloud Vertex AI Batch Prediction Job APIs."""
@@ -18,12 +18,10 @@
18
18
  """This module contains a Google Cloud Vertex AI hook."""
19
19
  from __future__ import annotations
20
20
 
21
- from typing import Sequence
21
+ from typing import TYPE_CHECKING, Sequence
22
22
 
23
23
  from google.api_core.client_options import ClientOptions
24
24
  from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
25
- from google.api_core.operation import Operation
26
- from google.api_core.retry import Retry
27
25
  from google.cloud.aiplatform import (
28
26
  CustomContainerTrainingJob,
29
27
  CustomPythonPackageTrainingJob,
@@ -32,17 +30,21 @@ from google.cloud.aiplatform import (
32
30
  models,
33
31
  )
34
32
  from google.cloud.aiplatform_v1 import JobServiceClient, PipelineServiceClient
35
- from google.cloud.aiplatform_v1.services.job_service.pagers import ListCustomJobsPager
36
- from google.cloud.aiplatform_v1.services.pipeline_service.pagers import (
37
- ListPipelineJobsPager,
38
- ListTrainingPipelinesPager,
39
- )
40
- from google.cloud.aiplatform_v1.types import CustomJob, PipelineJob, TrainingPipeline
41
33
 
42
34
  from airflow import AirflowException
43
35
  from airflow.providers.google.common.consts import CLIENT_INFO
44
36
  from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
45
37
 
38
+ if TYPE_CHECKING:
39
+ from google.api_core.operation import Operation
40
+ from google.api_core.retry import Retry
41
+ from google.cloud.aiplatform_v1.services.job_service.pagers import ListCustomJobsPager
42
+ from google.cloud.aiplatform_v1.services.pipeline_service.pagers import (
43
+ ListPipelineJobsPager,
44
+ ListTrainingPipelinesPager,
45
+ )
46
+ from google.cloud.aiplatform_v1.types import CustomJob, PipelineJob, TrainingPipeline
47
+
46
48
 
47
49
  class CustomJobHook(GoogleBaseHook):
48
50
  """Hook for Google Cloud Vertex AI Custom Job APIs."""
@@ -18,25 +18,27 @@
18
18
  """This module contains a Google Cloud Vertex AI hook."""
19
19
  from __future__ import annotations
20
20
 
21
- from typing import Sequence
21
+ from typing import TYPE_CHECKING, Sequence
22
22
 
23
23
  from google.api_core.client_options import ClientOptions
24
24
  from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
25
- from google.api_core.operation import Operation
26
- from google.api_core.retry import Retry
27
25
  from google.cloud.aiplatform_v1 import DatasetServiceClient
28
- from google.cloud.aiplatform_v1.services.dataset_service.pagers import (
29
- ListAnnotationsPager,
30
- ListDataItemsPager,
31
- ListDatasetsPager,
32
- )
33
- from google.cloud.aiplatform_v1.types import AnnotationSpec, Dataset, ExportDataConfig, ImportDataConfig
34
- from google.protobuf.field_mask_pb2 import FieldMask
35
26
 
36
27
  from airflow import AirflowException
37
28
  from airflow.providers.google.common.consts import CLIENT_INFO
38
29
  from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
39
30
 
31
+ if TYPE_CHECKING:
32
+ from google.api_core.operation import Operation
33
+ from google.api_core.retry import Retry
34
+ from google.cloud.aiplatform_v1.services.dataset_service.pagers import (
35
+ ListAnnotationsPager,
36
+ ListDataItemsPager,
37
+ ListDatasetsPager,
38
+ )
39
+ from google.cloud.aiplatform_v1.types import AnnotationSpec, Dataset, ExportDataConfig, ImportDataConfig
40
+ from google.protobuf.field_mask_pb2 import FieldMask
41
+
40
42
 
41
43
  class DatasetHook(GoogleBaseHook):
42
44
  """Hook for Google Cloud Vertex AI Dataset APIs."""
@@ -28,20 +28,22 @@
28
28
  """
29
29
  from __future__ import annotations
30
30
 
31
- from typing import Sequence
31
+ from typing import TYPE_CHECKING, Sequence
32
32
 
33
33
  from google.api_core.client_options import ClientOptions
34
34
  from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
35
- from google.api_core.operation import Operation
36
- from google.api_core.retry import Retry
37
35
  from google.cloud.aiplatform_v1 import EndpointServiceClient
38
- from google.cloud.aiplatform_v1.services.endpoint_service.pagers import ListEndpointsPager
39
- from google.cloud.aiplatform_v1.types import DeployedModel, Endpoint
40
- from google.protobuf.field_mask_pb2 import FieldMask
41
36
 
42
37
  from airflow import AirflowException
43
38
  from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
44
39
 
40
+ if TYPE_CHECKING:
41
+ from google.api_core.operation import Operation
42
+ from google.api_core.retry import Retry
43
+ from google.cloud.aiplatform_v1.services.endpoint_service.pagers import ListEndpointsPager
44
+ from google.cloud.aiplatform_v1.types import DeployedModel, Endpoint
45
+ from google.protobuf.field_mask_pb2 import FieldMask
46
+
45
47
 
46
48
  class EndpointServiceHook(GoogleBaseHook):
47
49
  """Hook for Google Cloud Vertex AI Endpoint Service APIs."""
@@ -27,19 +27,21 @@
27
27
  """
28
28
  from __future__ import annotations
29
29
 
30
- from typing import Sequence
30
+ from typing import TYPE_CHECKING, Sequence
31
31
 
32
32
  from google.api_core.client_options import ClientOptions
33
33
  from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
34
- from google.api_core.operation import Operation
35
- from google.api_core.retry import Retry
36
34
  from google.cloud.aiplatform import CustomJob, HyperparameterTuningJob, gapic, hyperparameter_tuning
37
35
  from google.cloud.aiplatform_v1 import JobServiceClient, types
38
- from google.cloud.aiplatform_v1.services.job_service.pagers import ListHyperparameterTuningJobsPager
39
36
 
40
37
  from airflow import AirflowException
41
38
  from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
42
39
 
40
+ if TYPE_CHECKING:
41
+ from google.api_core.operation import Operation
42
+ from google.api_core.retry import Retry
43
+ from google.cloud.aiplatform_v1.services.job_service.pagers import ListHyperparameterTuningJobsPager
44
+
43
45
 
44
46
  class HyperparameterTuningJobHook(GoogleBaseHook):
45
47
  """Hook for Google Cloud Vertex AI Hyperparameter Tuning Job APIs."""
@@ -24,19 +24,21 @@
24
24
  """
25
25
  from __future__ import annotations
26
26
 
27
- from typing import Sequence
27
+ from typing import TYPE_CHECKING, Sequence
28
28
 
29
29
  from google.api_core.client_options import ClientOptions
30
30
  from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
31
- from google.api_core.operation import Operation
32
- from google.api_core.retry import Retry
33
31
  from google.cloud.aiplatform_v1 import ModelServiceClient
34
- from google.cloud.aiplatform_v1.services.model_service.pagers import ListModelsPager
35
- from google.cloud.aiplatform_v1.types import Model, model_service
36
32
 
37
33
  from airflow import AirflowException
38
34
  from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
39
35
 
36
+ if TYPE_CHECKING:
37
+ from google.api_core.operation import Operation
38
+ from google.api_core.retry import Retry
39
+ from google.cloud.aiplatform_v1.services.model_service.pagers import ListModelsPager
40
+ from google.cloud.aiplatform_v1.types import Model, model_service
41
+
40
42
 
41
43
  class ModelServiceHook(GoogleBaseHook):
42
44
  """Hook for Google Cloud Vertex AI Endpoint Service APIs."""
@@ -18,11 +18,9 @@
18
18
  """This module contains a Google Cloud Video Intelligence Hook."""
19
19
  from __future__ import annotations
20
20
 
21
- from typing import Sequence
21
+ from typing import TYPE_CHECKING, Sequence
22
22
 
23
23
  from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
24
- from google.api_core.operation import Operation
25
- from google.api_core.retry import Retry
26
24
  from google.cloud.videointelligence_v1 import (
27
25
  Feature,
28
26
  VideoContext,
@@ -32,6 +30,10 @@ from google.cloud.videointelligence_v1 import (
32
30
  from airflow.providers.google.common.consts import CLIENT_INFO
33
31
  from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
34
32
 
33
+ if TYPE_CHECKING:
34
+ from google.api_core.operation import Operation
35
+ from google.api_core.retry import Retry
36
+
35
37
 
36
38
  class CloudVideoIntelligenceHook(GoogleBaseHook):
37
39
  """