Add kaggle user agent to gcp integration client.s

Author: vimota

patches/kaggle_gcp.py

@@ -2,13 +2,15 @@
 import inspect
 from google.auth import credentials
 from google.auth.exceptions import RefreshError
+from google.api_core.client_info import ClientInfo
 from google.cloud import bigquery
 from google.cloud.exceptions import Forbidden
 from google.cloud.bigquery._http import Connection
 from kaggle_secrets import GcpTarget, UserSecretsClient
 
 from log import Log
 
+KAGGLE_GCP_CLIENT_USER_AGENT="kaggle-gcp-client/1.0"
 
 def get_integrations():
     kernel_integrations_var = os.getenv("KAGGLE_KERNEL_INTEGRATIONS")
@@ -165,6 +167,7 @@ def monkeypatch_bq(bq_client, *args, **kwargs):
                 Log.info("No project specified while using the unmodified client.")
                 print('Please ensure you specify a project id when creating the client'
                     ' in order to use your BigQuery account.')
+            kwargs['client_info'] = set_kaggle_user_agent(kwargs.get('client_info'))
             return bq_client(*args, **kwargs)
 
     # Monkey patches BigQuery client creation to use proxy or user-connected GCP account.
@@ -183,11 +186,28 @@ def patched_init(self, *args, **kwargs):
         if specified_credentials is None:
             Log.info("No credentials specified, using KaggleKernelCredentials.")
             kwargs['credentials'] = kaggle_kernel_credentials
+        
+        # TODO(vimota): Remove the exclusion of TablesClient once
+        # the client has fixed the error:
+        # `multiple values for keyword argument 'client_info'``
+        from google.cloud import automl_v1beta1
+        if (client_klass != automl_v1beta1.TablesClient):
+            kwargs['client_info'] = set_kaggle_user_agent(kwargs.get('client_info'))
+
+
         return client_init(self, *args, **kwargs)
 
     if (not has_been_monkeypatched(client_klass.__init__)):
         client_klass.__init__ = patched_init
 
+def set_kaggle_user_agent(client_info: ClientInfo):
+    # Add kaggle client user agent in order to attribute usage.
+    if client_info is None:
+        client_info = ClientInfo(user_agent=KAGGLE_GCP_CLIENT_USER_AGENT)
+    else:
+        client_info.user_agent = KAGGLE_GCP_CLIENT_USER_AGENT
+    return client_info
+
 def init_gcs():
     is_user_secrets_token_set = "KAGGLE_USER_SECRETS_TOKEN" in os.environ
     from google.cloud import storage

tests/test_automl.py

@@ -19,9 +19,15 @@ def test_version(self):
         self.assertGreaterEqual(version, 0.5);
 
     class FakeClient:
-        def __init__(self, credentials=None, **kwargs):
+        def __init__(self, credentials=None, client_info=None, **kwargs):
             self.credentials = credentials
 
+            class FakeConnection():
+                def __init__(self, user_agent):
+                    self.user_agent = user_agent
+            if (client_info is not None):
+                self._connection = FakeConnection(client_info.user_agent)
+
     @patch("google.cloud.automl_v1beta1.AutoMlClient", new=FakeClient)
     def test_user_provided_credentials(self):
         credentials = _make_credentials()
@@ -64,6 +70,7 @@ def test_default_credentials_automl_client(self):
             automl_client = automl.AutoMlClient()
             self.assertIsNotNone(automl_client.credentials)
             self.assertIsInstance(automl_client.credentials, KaggleKernelCredentials)
+            self.assertTrue(automl_client._connection.user_agent.startswith("kaggle-gcp-client/1.0"))
 
     @patch("google.cloud.automl_v1beta1.TablesClient", new=FakeClient)
     def test_default_credentials_tables_client(self):

tests/test_bigquery.py

@@ -103,6 +103,7 @@ def test_project_with_connected_account_default_credentials(self, mock_access_to
         env.set('KAGGLE_KERNEL_INTEGRATIONS', 'BIGQUERY')
         with env:
             client = bigquery.Client(project='ANOTHER_PROJECT')
+            self.assertTrue(client._connection.user_agent.startswith("kaggle-gcp-client/1.0"))
             self._test_integration(client)
 
     @patch.object(Connection, 'API_BASE_URL')

tests/test_gcs.py

@@ -44,6 +44,7 @@ def test_default_credentials_gcs_enabled(self):
             init_gcs()
             client = storage.Client(project="xyz")
             self.assertIsInstance(client._credentials, KaggleKernelCredentials)
+            self.assertTrue(client._connection.user_agent.startswith("kaggle-gcp-client/1.0"))
 
     def test_monkeypatching_idempotent(self):
         env = EnvironmentVarGuard()