Extract ClientInfo to module level (#21554)

apache · Feb 15, 2022 · 1b568d7 · 1b568d7
1 parent 7cca824
commit 1b568d7
Show file tree

Hide file tree

Showing 51 changed files with 184 additions and 286 deletions.
diff --git a/airflow/providers/google/cloud/_internal_client/secret_manager_client.py b/airflow/providers/google/cloud/_internal_client/secret_manager_client.py
@@ -21,17 +21,17 @@
 
 import google
 
+from airflow.providers.google.common.consts import CLIENT_INFO
+
 if sys.version_info >= (3, 8):
     from functools import cached_property
 else:
     from cached_property import cached_property
 
 from google.api_core.exceptions import InvalidArgument, NotFound, PermissionDenied
-from google.api_core.gapic_v1.client_info import ClientInfo
 from google.cloud.secretmanager_v1 import SecretManagerServiceClient
 
 from airflow.utils.log.logging_mixin import LoggingMixin
-from airflow.version import version
 
 SECRET_ID_PATTERN = r"^[a-zA-Z0-9-_]*$"
 
@@ -65,9 +65,7 @@ def is_valid_secret_name(secret_name: str) -> bool:
     @cached_property
     def client(self) -> SecretManagerServiceClient:
         """Create an authenticated KMS client"""
-        _client = SecretManagerServiceClient(
-            credentials=self.credentials, client_info=ClientInfo(client_library_version='airflow_v' + version)
-        )
+        _client = SecretManagerServiceClient(credentials=self.credentials, client_info=CLIENT_INFO)
         return _client
 
     def get_secret(self, secret_id: str, project_id: str, secret_version: str = 'latest') -> Optional[str]:

diff --git a/airflow/providers/google/cloud/hooks/automl.py b/airflow/providers/google/cloud/hooks/automl.py
@@ -32,6 +32,8 @@
     ListTableSpecsPager,
 )
 
+from airflow.providers.google.common.consts import CLIENT_INFO
+
 if sys.version_info >= (3, 8):
     from functools import cached_property
 else:
@@ -90,7 +92,7 @@ def get_conn(self) -> AutoMlClient:
         :rtype: google.cloud.automl_v1beta1.AutoMlClient
         """
         if self._client is None:
-            self._client = AutoMlClient(credentials=self._get_credentials(), client_info=self.client_info)
+            self._client = AutoMlClient(credentials=self._get_credentials(), client_info=CLIENT_INFO)
         return self._client
 
     @cached_property
@@ -101,7 +103,7 @@ def prediction_client(self) -> PredictionServiceClient:
         :return: Google Cloud AutoML PredictionServiceClient client object.
         :rtype: google.cloud.automl_v1beta1.PredictionServiceClient
         """
-        return PredictionServiceClient(credentials=self._get_credentials(), client_info=self.client_info)
+        return PredictionServiceClient(credentials=self._get_credentials(), client_info=CLIENT_INFO)
 
     @GoogleBaseHook.fallback_to_default_project_id
     def create_model(

diff --git a/airflow/providers/google/cloud/hooks/bigquery.py b/airflow/providers/google/cloud/hooks/bigquery.py
@@ -55,6 +55,7 @@
 
 from airflow.exceptions import AirflowException
 from airflow.hooks.dbapi import DbApiHook
+from airflow.providers.google.common.consts import CLIENT_INFO
 from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
 from airflow.utils.helpers import convert_camel_to_snake
 from airflow.utils.log.logging_mixin import LoggingMixin
@@ -146,7 +147,7 @@ def get_client(self, project_id: Optional[str] = None, location: Optional[str] =
         :return:
         """
         return Client(
-            client_info=self.client_info,
+            client_info=CLIENT_INFO,
             project=project_id,
             location=location,
             credentials=self._get_credentials(),

diff --git a/airflow/providers/google/cloud/hooks/bigquery_dts.py b/airflow/providers/google/cloud/hooks/bigquery_dts.py
@@ -29,6 +29,7 @@
 )
 from googleapiclient.discovery import Resource
 
+from airflow.providers.google.common.consts import CLIENT_INFO
 from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID, GoogleBaseHook
 
 
@@ -95,7 +96,7 @@ def get_conn(self) -> DataTransferServiceClient:
         """
         if not self._conn:
             self._conn = DataTransferServiceClient(
-                credentials=self._get_credentials(), client_info=self.client_info
+                credentials=self._get_credentials(), client_info=CLIENT_INFO
             )
         return self._conn
 

diff --git a/airflow/providers/google/cloud/hooks/bigtable.py b/airflow/providers/google/cloud/hooks/bigtable.py
@@ -27,6 +27,7 @@
 from google.cloud.bigtable.table import ClusterState, Table
 from google.cloud.bigtable_admin_v2 import enums
 
+from airflow.providers.google.common.consts import CLIENT_INFO
 from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
 
 
@@ -56,7 +57,7 @@ def _get_client(self, project_id: str):
             self._client = Client(
                 project=project_id,
                 credentials=self._get_credentials(),
-                client_info=self.client_info,
+                client_info=CLIENT_INFO,
                 admin=True,
             )
         return self._client

diff --git a/airflow/providers/google/cloud/hooks/cloud_build.py b/airflow/providers/google/cloud/hooks/cloud_build.py
@@ -26,6 +26,7 @@
 from google.cloud.devtools.cloudbuild_v1.types import Build, BuildTrigger, RepoSource
 
 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
 
 # Time to sleep between active checks of the operation results
@@ -84,7 +85,7 @@ def get_conn(self) -> CloudBuildClient:
         :rtype: `google.cloud.devtools.cloudbuild_v1.CloudBuildClient`
         """
         if not self._client:
-            self._client = CloudBuildClient(credentials=self._get_credentials(), client_info=self.client_info)
+            self._client = CloudBuildClient(credentials=self._get_credentials(), client_info=CLIENT_INFO)
         return self._client
 
     @GoogleBaseHook.fallback_to_default_project_id

diff --git a/airflow/providers/google/cloud/hooks/cloud_composer.py b/airflow/providers/google/cloud/hooks/cloud_composer.py
@@ -29,6 +29,7 @@
 from google.protobuf.field_mask_pb2 import FieldMask
 
 from airflow import AirflowException
+from airflow.providers.google.common.consts import CLIENT_INFO
 from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
 
 
@@ -41,7 +42,7 @@ def get_environment_client(self) -> EnvironmentsClient:
         """Retrieves client library object that allow access Environments service."""
         return EnvironmentsClient(
             credentials=self._get_credentials(),
-            client_info=self.client_info,
+            client_info=CLIENT_INFO,
             client_options=self.client_options,
         )
 
@@ -51,7 +52,7 @@ def get_image_versions_client(
         """Retrieves client library object that allow access Image Versions service."""
         return ImageVersionsClient(
             credentials=self._get_credentials(),
-            client_info=self.client_info,
+            client_info=CLIENT_INFO,
             client_options=self.client_options,
         )
 

diff --git a/airflow/providers/google/cloud/hooks/datacatalog.py b/airflow/providers/google/cloud/hooks/datacatalog.py
@@ -32,6 +32,7 @@
 from google.protobuf.field_mask_pb2 import FieldMask
 
 from airflow import AirflowException
+from airflow.providers.google.common.consts import CLIENT_INFO
 from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID, GoogleBaseHook
 
 
@@ -69,9 +70,7 @@ def __init__(
     def get_conn(self) -> DataCatalogClient:
         """Retrieves client library object that allow access to Cloud Data Catalog service."""
         if not self._client:
-            self._client = DataCatalogClient(
-                credentials=self._get_credentials(), client_info=self.client_info
-            )
+            self._client = DataCatalogClient(credentials=self._get_credentials(), client_info=CLIENT_INFO)
         return self._client
 
     @GoogleBaseHook.fallback_to_default_project_id

diff --git a/airflow/providers/google/cloud/hooks/dataproc.py b/airflow/providers/google/cloud/hooks/dataproc.py
@@ -41,6 +41,7 @@
 from google.protobuf.field_mask_pb2 import FieldMask
 
 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 airflow.version import version as airflow_version
 
@@ -214,7 +215,7 @@ def get_cluster_client(
             client_options = {'api_endpoint': f'{region}-dataproc.googleapis.com:443'}
 
         return ClusterControllerClient(
-            credentials=self._get_credentials(), client_info=self.client_info, client_options=client_options
+            credentials=self._get_credentials(), client_info=CLIENT_INFO, client_options=client_options
         )
 
     def get_template_client(
@@ -234,7 +235,7 @@ def get_template_client(
             client_options = {'api_endpoint': f'{region}-dataproc.googleapis.com:443'}
 
         return WorkflowTemplateServiceClient(
-            credentials=self._get_credentials(), client_info=self.client_info, client_options=client_options
+            credentials=self._get_credentials(), client_info=CLIENT_INFO, client_options=client_options
         )
 
     def get_job_client(
@@ -254,7 +255,7 @@ def get_job_client(
             client_options = {'api_endpoint': f'{region}-dataproc.googleapis.com:443'}
 
         return JobControllerClient(
-            credentials=self._get_credentials(), client_info=self.client_info, client_options=client_options
+            credentials=self._get_credentials(), client_info=CLIENT_INFO, client_options=client_options
         )
 
     def get_batch_client(
@@ -274,7 +275,7 @@ def get_batch_client(
             client_options = {'api_endpoint': f'{region}-dataproc.googleapis.com:443'}
 
         return BatchControllerClient(
-            credentials=self._get_credentials(), client_info=self.client_info, client_options=client_options
+            credentials=self._get_credentials(), client_info=CLIENT_INFO, client_options=client_options
         )
 
     def wait_for_operation(self, operation: Operation, timeout: Optional[float] = None):

diff --git a/airflow/providers/google/cloud/hooks/dataproc_metastore.py b/airflow/providers/google/cloud/hooks/dataproc_metastore.py
@@ -28,6 +28,7 @@
 from google.protobuf.field_mask_pb2 import FieldMask
 
 from airflow.exceptions import AirflowException
+from airflow.providers.google.common.consts import CLIENT_INFO
 from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
 
 
@@ -39,7 +40,7 @@ def get_dataproc_metastore_client(self) -> DataprocMetastoreClient:
         client_options = {'api_endpoint': 'metastore.googleapis.com:443'}
 
         return DataprocMetastoreClient(
-            credentials=self._get_credentials(), client_info=self.client_info, client_options=client_options
+            credentials=self._get_credentials(), client_info=CLIENT_INFO, client_options=client_options
         )
 
     def wait_for_operation(self, timeout: Optional[float], operation: Operation):

diff --git a/airflow/providers/google/cloud/hooks/dlp.py b/airflow/providers/google/cloud/hooks/dlp.py
@@ -55,6 +55,7 @@
 )
 
 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
 
 DLP_JOB_PATH_PATTERN = "^projects/[^/]+/dlpJobs/(?P<job>.*?)$"
@@ -104,7 +105,7 @@ def get_conn(self) -> DlpServiceClient:
         :rtype: google.cloud.dlp_v2.DlpServiceClient
         """
         if not self._client:
-            self._client = DlpServiceClient(credentials=self._get_credentials(), client_info=self.client_info)
+            self._client = DlpServiceClient(credentials=self._get_credentials(), client_info=CLIENT_INFO)
         return self._client
 
     @GoogleBaseHook.fallback_to_default_project_id

diff --git a/airflow/providers/google/cloud/hooks/gcs.py b/airflow/providers/google/cloud/hooks/gcs.py
@@ -40,6 +40,7 @@
 
 from airflow.exceptions import AirflowException
 from airflow.providers.google.cloud.utils.helpers import normalize_directory_path
+from airflow.providers.google.common.consts import CLIENT_INFO
 from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
 from airflow.utils import timezone
 from airflow.version import version
@@ -155,7 +156,7 @@ def get_conn(self) -> storage.Client:
         """Returns a Google Cloud Storage service object."""
         if not self._conn:
             self._conn = storage.Client(
-                credentials=self._get_credentials(), client_info=self.client_info, project=self.project_id
+                credentials=self._get_credentials(), client_info=CLIENT_INFO, project=self.project_id
             )
 
         return self._conn

diff --git a/airflow/providers/google/cloud/hooks/kms.py b/airflow/providers/google/cloud/hooks/kms.py
@@ -25,6 +25,7 @@
 from google.api_core.retry import Retry
 from google.cloud.kms_v1 import KeyManagementServiceClient
 
+from airflow.providers.google.common.consts import CLIENT_INFO
 from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
 
 
@@ -78,7 +79,7 @@ def get_conn(self) -> KeyManagementServiceClient:
         """
         if not self._conn:
             self._conn = KeyManagementServiceClient(
-                credentials=self._get_credentials(), client_info=self.client_info
+                credentials=self._get_credentials(), client_info=CLIENT_INFO
             )
         return self._conn
 

diff --git a/airflow/providers/google/cloud/hooks/kubernetes_engine.py b/airflow/providers/google/cloud/hooks/kubernetes_engine.py
@@ -41,6 +41,7 @@
 
 from airflow import version
 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
 
 OPERATIONAL_POLL_INTERVAL = 15
@@ -77,9 +78,7 @@ def get_conn(self) -> container_v1.ClusterManagerClient:
         """
         if self._client is None:
             credentials = self._get_credentials()
-            self._client = container_v1.ClusterManagerClient(
-                credentials=credentials, client_info=self.client_info
-            )
+            self._client = container_v1.ClusterManagerClient(credentials=credentials, client_info=CLIENT_INFO)
         return self._client
 
     # To preserve backward compatibility

diff --git a/airflow/providers/google/cloud/hooks/natural_language.py b/airflow/providers/google/cloud/hooks/natural_language.py
@@ -31,6 +31,7 @@
     Document,
 )
 
+from airflow.providers.google.common.consts import CLIENT_INFO
 from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
 
 
@@ -73,9 +74,7 @@ def get_conn(self) -> LanguageServiceClient:
         :rtype: google.cloud.language_v1.LanguageServiceClient
         """
         if not self._conn:
-            self._conn = LanguageServiceClient(
-                credentials=self._get_credentials(), client_info=self.client_info
-            )
+            self._conn = LanguageServiceClient(credentials=self._get_credentials(), client_info=CLIENT_INFO)
         return self._conn
 
     @GoogleBaseHook.quota_retry()

diff --git a/airflow/providers/google/cloud/hooks/os_login.py b/airflow/providers/google/cloud/hooks/os_login.py
@@ -27,6 +27,7 @@
 from google.api_core.retry import Retry
 from google.cloud.oslogin_v1 import ImportSshPublicKeyResponse, OsLoginServiceClient
 
+from airflow.providers.google.common.consts import CLIENT_INFO
 from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID, GoogleBaseHook
 
 
@@ -56,7 +57,7 @@ def get_conn(self) -> OsLoginServiceClient:
         if self._conn:
             return self._conn
 
-        self._conn = OsLoginServiceClient(credentials=self._get_credentials(), client_info=self.client_info)
+        self._conn = OsLoginServiceClient(credentials=self._get_credentials(), client_info=CLIENT_INFO)
         return self._conn
 
     @GoogleBaseHook.fallback_to_default_project_id

diff --git a/airflow/providers/google/cloud/hooks/pubsub.py b/airflow/providers/google/cloud/hooks/pubsub.py
@@ -29,6 +29,8 @@
 from typing import Dict, List, Optional, Sequence, Tuple, Union
 from uuid import uuid4
 
+from airflow.providers.google.common.consts import CLIENT_INFO
+
 if sys.version_info >= (3, 8):
     from functools import cached_property
 else:
@@ -86,7 +88,7 @@ def get_conn(self) -> PublisherClient:
         :rtype: google.cloud.pubsub_v1.PublisherClient
         """
         if not self._client:
-            self._client = PublisherClient(credentials=self._get_credentials(), client_info=self.client_info)
+            self._client = PublisherClient(credentials=self._get_credentials(), client_info=CLIENT_INFO)
         return self._client
 
     @cached_property
@@ -97,7 +99,7 @@ def subscriber_client(self) -> SubscriberClient:
         :return: Google Cloud Pub/Sub client object.
         :rtype: google.cloud.pubsub_v1.SubscriberClient
         """
-        return SubscriberClient(credentials=self._get_credentials(), client_info=self.client_info)
+        return SubscriberClient(credentials=self._get_credentials(), client_info=CLIENT_INFO)
 
     @GoogleBaseHook.fallback_to_default_project_id
     def publish(

diff --git a/airflow/providers/google/cloud/hooks/spanner.py b/airflow/providers/google/cloud/hooks/spanner.py
@@ -26,6 +26,7 @@
 from google.longrunning.operations_grpc_pb2 import Operation
 
 from airflow.exceptions import AirflowException
+from airflow.providers.google.common.consts import CLIENT_INFO
 from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
 
 
@@ -60,7 +61,7 @@ def _get_client(self, project_id: str) -> Client:
         """
         if not self._client:
             self._client = Client(
-                project=project_id, credentials=self._get_credentials(), client_info=self.client_info
+                project=project_id, credentials=self._get_credentials(), client_info=CLIENT_INFO
             )
         return self._client
 

diff --git a/airflow/providers/google/cloud/hooks/speech_to_text.py b/airflow/providers/google/cloud/hooks/speech_to_text.py
@@ -22,6 +22,7 @@
 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
 
 
@@ -64,7 +65,7 @@ def get_conn(self) -> SpeechClient:
         :rtype: google.cloud.speech_v1.SpeechClient
         """
         if not self._client:
-            self._client = SpeechClient(credentials=self._get_credentials(), client_info=self.client_info)
+            self._client = SpeechClient(credentials=self._get_credentials(), client_info=CLIENT_INFO)
         return self._client
 
     @GoogleBaseHook.quota_retry()