From 781be08d7b63373af5c81208bc793850ef23124c Mon Sep 17 00:00:00 2001
From: Victor Mota <vimota@gmail.com>
Date: Sat, 7 Dec 2019 08:34:51 +0000
Subject: [PATCH] Add kaggle user agent to gcp integration client.s

---
 patches/kaggle_gcp.py  | 20 ++++++++++++++++++++
 tests/test_automl.py   |  9 ++++++++-
 tests/test_bigquery.py |  1 +
 tests/test_gcs.py      |  1 +
 4 files changed, 30 insertions(+), 1 deletion(-)

diff --git a/patches/kaggle_gcp.py b/patches/kaggle_gcp.py
index aa9b4309..01497822 100644
--- a/patches/kaggle_gcp.py
+++ b/patches/kaggle_gcp.py
@@ -2,6 +2,7 @@
 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
@@ -9,6 +10,7 @@
 
 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
diff --git a/tests/test_automl.py b/tests/test_automl.py
index 1cb1e110..367f5ee7 100644
--- a/tests/test_automl.py
+++ b/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):
diff --git a/tests/test_bigquery.py b/tests/test_bigquery.py
index a71f9545..4b152b78 100644
--- a/tests/test_bigquery.py
+++ b/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')
diff --git a/tests/test_gcs.py b/tests/test_gcs.py
index c6b2c45e..9902a877 100644
--- a/tests/test_gcs.py
+++ b/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()