Skip to content

Commit 1b568d7

Browse files
pierrejeambrunpierrejeambrun
authored
Extract ClientInfo to module level (#21554)
1 parent 7cca824 commit 1b568d7

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

51 files changed

+184
-286
lines changed

β€Žairflow/providers/google/cloud/_internal_client/secret_manager_client.py

Lines changed: 3 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -21,17 +21,17 @@
2121

2222
import google
2323

24+
from airflow.providers.google.common.consts import CLIENT_INFO
25+
2426
if sys.version_info >= (3, 8):
2527
from functools import cached_property
2628
else:
2729
from cached_property import cached_property
2830

2931
from google.api_core.exceptions import InvalidArgument, NotFound, PermissionDenied
30-
from google.api_core.gapic_v1.client_info import ClientInfo
3132
from google.cloud.secretmanager_v1 import SecretManagerServiceClient
3233

3334
from airflow.utils.log.logging_mixin import LoggingMixin
34-
from airflow.version import version
3535

3636
SECRET_ID_PATTERN = r"^[a-zA-Z0-9-_]*$"
3737

@@ -65,9 +65,7 @@ def is_valid_secret_name(secret_name: str) -> bool:
6565
@cached_property
6666
def client(self) -> SecretManagerServiceClient:
6767
"""Create an authenticated KMS client"""
68-
_client = SecretManagerServiceClient(
69-
credentials=self.credentials, client_info=ClientInfo(client_library_version='airflow_v' + version)
70-
)
68+
_client = SecretManagerServiceClient(credentials=self.credentials, client_info=CLIENT_INFO)
7169
return _client
7270

7371
def get_secret(self, secret_id: str, project_id: str, secret_version: str = 'latest') -> Optional[str]:

β€Žairflow/providers/google/cloud/hooks/automl.py

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,8 @@
3232
ListTableSpecsPager,
3333
)
3434

35+
from airflow.providers.google.common.consts import CLIENT_INFO
36+
3537
if sys.version_info >= (3, 8):
3638
from functools import cached_property
3739
else:
@@ -90,7 +92,7 @@ def get_conn(self) -> AutoMlClient:
9092
:rtype: google.cloud.automl_v1beta1.AutoMlClient
9193
"""
9294
if self._client is None:
93-
self._client = AutoMlClient(credentials=self._get_credentials(), client_info=self.client_info)
95+
self._client = AutoMlClient(credentials=self._get_credentials(), client_info=CLIENT_INFO)
9496
return self._client
9597

9698
@cached_property
@@ -101,7 +103,7 @@ def prediction_client(self) -> PredictionServiceClient:
101103
:return: Google Cloud AutoML PredictionServiceClient client object.
102104
:rtype: google.cloud.automl_v1beta1.PredictionServiceClient
103105
"""
104-
return PredictionServiceClient(credentials=self._get_credentials(), client_info=self.client_info)
106+
return PredictionServiceClient(credentials=self._get_credentials(), client_info=CLIENT_INFO)
105107

106108
@GoogleBaseHook.fallback_to_default_project_id
107109
def create_model(

β€Žairflow/providers/google/cloud/hooks/bigquery.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -55,6 +55,7 @@
5555

5656
from airflow.exceptions import AirflowException
5757
from airflow.hooks.dbapi import DbApiHook
58+
from airflow.providers.google.common.consts import CLIENT_INFO
5859
from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
5960
from airflow.utils.helpers import convert_camel_to_snake
6061
from airflow.utils.log.logging_mixin import LoggingMixin
@@ -146,7 +147,7 @@ def get_client(self, project_id: Optional[str] = None, location: Optional[str] =
146147
:return:
147148
"""
148149
return Client(
149-
client_info=self.client_info,
150+
client_info=CLIENT_INFO,
150151
project=project_id,
151152
location=location,
152153
credentials=self._get_credentials(),

β€Žairflow/providers/google/cloud/hooks/bigquery_dts.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@
2929
)
3030
from googleapiclient.discovery import Resource
3131

32+
from airflow.providers.google.common.consts import CLIENT_INFO
3233
from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID, GoogleBaseHook
3334

3435

@@ -95,7 +96,7 @@ def get_conn(self) -> DataTransferServiceClient:
9596
"""
9697
if not self._conn:
9798
self._conn = DataTransferServiceClient(
98-
credentials=self._get_credentials(), client_info=self.client_info
99+
credentials=self._get_credentials(), client_info=CLIENT_INFO
99100
)
100101
return self._conn
101102

β€Žairflow/providers/google/cloud/hooks/bigtable.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@
2727
from google.cloud.bigtable.table import ClusterState, Table
2828
from google.cloud.bigtable_admin_v2 import enums
2929

30+
from airflow.providers.google.common.consts import CLIENT_INFO
3031
from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
3132

3233

@@ -56,7 +57,7 @@ def _get_client(self, project_id: str):
5657
self._client = Client(
5758
project=project_id,
5859
credentials=self._get_credentials(),
59-
client_info=self.client_info,
60+
client_info=CLIENT_INFO,
6061
admin=True,
6162
)
6263
return self._client

β€Žairflow/providers/google/cloud/hooks/cloud_build.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@
2626
from google.cloud.devtools.cloudbuild_v1.types import Build, BuildTrigger, RepoSource
2727

2828
from airflow.exceptions import AirflowException
29+
from airflow.providers.google.common.consts import CLIENT_INFO
2930
from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID, GoogleBaseHook
3031

3132
# Time to sleep between active checks of the operation results
@@ -84,7 +85,7 @@ def get_conn(self) -> CloudBuildClient:
8485
:rtype: `google.cloud.devtools.cloudbuild_v1.CloudBuildClient`
8586
"""
8687
if not self._client:
87-
self._client = CloudBuildClient(credentials=self._get_credentials(), client_info=self.client_info)
88+
self._client = CloudBuildClient(credentials=self._get_credentials(), client_info=CLIENT_INFO)
8889
return self._client
8990

9091
@GoogleBaseHook.fallback_to_default_project_id

β€Žairflow/providers/google/cloud/hooks/cloud_composer.py

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@
2929
from google.protobuf.field_mask_pb2 import FieldMask
3030

3131
from airflow import AirflowException
32+
from airflow.providers.google.common.consts import CLIENT_INFO
3233
from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
3334

3435

@@ -41,7 +42,7 @@ def get_environment_client(self) -> EnvironmentsClient:
4142
"""Retrieves client library object that allow access Environments service."""
4243
return EnvironmentsClient(
4344
credentials=self._get_credentials(),
44-
client_info=self.client_info,
45+
client_info=CLIENT_INFO,
4546
client_options=self.client_options,
4647
)
4748

@@ -51,7 +52,7 @@ def get_image_versions_client(
5152
"""Retrieves client library object that allow access Image Versions service."""
5253
return ImageVersionsClient(
5354
credentials=self._get_credentials(),
54-
client_info=self.client_info,
55+
client_info=CLIENT_INFO,
5556
client_options=self.client_options,
5657
)
5758

β€Žairflow/providers/google/cloud/hooks/datacatalog.py

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@
3232
from google.protobuf.field_mask_pb2 import FieldMask
3333

3434
from airflow import AirflowException
35+
from airflow.providers.google.common.consts import CLIENT_INFO
3536
from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID, GoogleBaseHook
3637

3738

@@ -69,9 +70,7 @@ def __init__(
6970
def get_conn(self) -> DataCatalogClient:
7071
"""Retrieves client library object that allow access to Cloud Data Catalog service."""
7172
if not self._client:
72-
self._client = DataCatalogClient(
73-
credentials=self._get_credentials(), client_info=self.client_info
74-
)
73+
self._client = DataCatalogClient(credentials=self._get_credentials(), client_info=CLIENT_INFO)
7574
return self._client
7675

7776
@GoogleBaseHook.fallback_to_default_project_id

β€Žairflow/providers/google/cloud/hooks/dataproc.py

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,7 @@
4141
from google.protobuf.field_mask_pb2 import FieldMask
4242

4343
from airflow.exceptions import AirflowException
44+
from airflow.providers.google.common.consts import CLIENT_INFO
4445
from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
4546
from airflow.version import version as airflow_version
4647

@@ -214,7 +215,7 @@ def get_cluster_client(
214215
client_options = {'api_endpoint': f'{region}-dataproc.googleapis.com:443'}
215216

216217
return ClusterControllerClient(
217-
credentials=self._get_credentials(), client_info=self.client_info, client_options=client_options
218+
credentials=self._get_credentials(), client_info=CLIENT_INFO, client_options=client_options
218219
)
219220

220221
def get_template_client(
@@ -234,7 +235,7 @@ def get_template_client(
234235
client_options = {'api_endpoint': f'{region}-dataproc.googleapis.com:443'}
235236

236237
return WorkflowTemplateServiceClient(
237-
credentials=self._get_credentials(), client_info=self.client_info, client_options=client_options
238+
credentials=self._get_credentials(), client_info=CLIENT_INFO, client_options=client_options
238239
)
239240

240241
def get_job_client(
@@ -254,7 +255,7 @@ def get_job_client(
254255
client_options = {'api_endpoint': f'{region}-dataproc.googleapis.com:443'}
255256

256257
return JobControllerClient(
257-
credentials=self._get_credentials(), client_info=self.client_info, client_options=client_options
258+
credentials=self._get_credentials(), client_info=CLIENT_INFO, client_options=client_options
258259
)
259260

260261
def get_batch_client(
@@ -274,7 +275,7 @@ def get_batch_client(
274275
client_options = {'api_endpoint': f'{region}-dataproc.googleapis.com:443'}
275276

276277
return BatchControllerClient(
277-
credentials=self._get_credentials(), client_info=self.client_info, client_options=client_options
278+
credentials=self._get_credentials(), client_info=CLIENT_INFO, client_options=client_options
278279
)
279280

280281
def wait_for_operation(self, operation: Operation, timeout: Optional[float] = None):

β€Žairflow/providers/google/cloud/hooks/dataproc_metastore.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@
2828
from google.protobuf.field_mask_pb2 import FieldMask
2929

3030
from airflow.exceptions import AirflowException
31+
from airflow.providers.google.common.consts import CLIENT_INFO
3132
from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
3233

3334

@@ -39,7 +40,7 @@ def get_dataproc_metastore_client(self) -> DataprocMetastoreClient:
3940
client_options = {'api_endpoint': 'metastore.googleapis.com:443'}
4041

4142
return DataprocMetastoreClient(
42-
credentials=self._get_credentials(), client_info=self.client_info, client_options=client_options
43+
credentials=self._get_credentials(), client_info=CLIENT_INFO, client_options=client_options
4344
)
4445

4546
def wait_for_operation(self, timeout: Optional[float], operation: Operation):

0 commit comments

Comments
 (0)