Patch the AutoML v1 clients to get Kaggle managed credentials (Kaggle#715)

When the integration initially launched, there was only the v1beta1 version of the AutoML clients. As of Nov. 2019, there is a v1 version but we never updated the code to patch the v1 clients. Addresses b/149264883

Author: mcollins42

patches/kaggle_gcp.py

@@ -186,14 +186,8 @@ 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'))
 
+        kwargs['client_info'] = set_kaggle_user_agent(kwargs.get('client_info'))
 
         return client_init(self, *args, **kwargs)
 
@@ -227,30 +221,35 @@ def init_gcs():
 
 def init_automl():
     is_user_secrets_token_set = "KAGGLE_USER_SECRETS_TOKEN" in os.environ
-    from google.cloud import automl_v1beta1 as automl
+    from google.cloud import automl, automl_v1beta1
     if not is_user_secrets_token_set:
-        return automl
+        return
 
     from kaggle_gcp import get_integrations
     if not get_integrations().has_automl():
-        return automl
+        return
 
     from kaggle_secrets import GcpTarget
     from kaggle_gcp import KaggleKernelCredentials
     kaggle_kernel_credentials = KaggleKernelCredentials(target=GcpTarget.AUTOML)
 
-    # The AutoML client library exposes 4 different client classes (AutoMlClient,
-    # TablesClient, PredictionServiceClient and GcsClient), so patch each of them.
-    # The same KaggleKernelCredentials are passed to all of them.
+    # Patch the 2 GA clients: AutoMlClient and PreditionServiceClient
     monkeypatch_client(automl.AutoMlClient, kaggle_kernel_credentials)
-    monkeypatch_client(automl.TablesClient, kaggle_kernel_credentials)
     monkeypatch_client(automl.PredictionServiceClient, kaggle_kernel_credentials)
-    # TODO(markcollins): The GcsClient in the AutoML client library version
-    # 0.5.0 doesn't handle credentials properly. I wrote PR:
-    # https://github.com/googleapis/google-cloud-python/pull/9299
-    # to address this issue. Add patching for GcsClient when we get a version of
-    # the library that includes the fixes.
-    return automl
+
+    # The AutoML client library exposes 3 different client classes (AutoMlClient,
+    # TablesClient, PredictionServiceClient), so patch each of them.
+    # The same KaggleKernelCredentials are passed to all of them.
+    # The GcsClient class is only used internally by TablesClient.
+
+    # The beta version of the clients that are now GA are included here for now.
+    # They are deprecated and will be removed by 1 May 2020.
+    monkeypatch_client(automl_v1beta1.AutoMlClient, kaggle_kernel_credentials)
+    monkeypatch_client(automl_v1beta1.PredictionServiceClient, kaggle_kernel_credentials)
+
+    # The TablesClient is still in beta, so this will not be deprecated until
+    # the TablesClient is GA.
+    monkeypatch_client(automl_v1beta1.TablesClient, kaggle_kernel_credentials)
 
 def init():
     init_bigquery()

tests/test_automl.py

@@ -4,19 +4,14 @@
 
 from kaggle_gcp import KaggleKernelCredentials, init_automl
 from test.support import EnvironmentVarGuard
-from google.cloud import storage, automl_v1beta1 as automl
-from packaging.version import parse
+from google.cloud import storage, automl_v1beta1, automl
 
 def _make_credentials():
     import google.auth.credentials
     return Mock(spec=google.auth.credentials.Credentials)
 
 class TestAutoMl(unittest.TestCase):
 
-    def test_version(self):
-        self.assertIsNotNone(automl.auto_ml_client._GAPIC_LIBRARY_VERSION)
-        self.assertGreaterEqual(parse(automl.auto_ml_client._GAPIC_LIBRARY_VERSION), parse("0.5.0"))
-
     class FakeClient:
         def __init__(self, credentials=None, client_info=None, **kwargs):
             self.credentials = credentials
@@ -27,7 +22,7 @@ def __init__(self, user_agent):
             if (client_info is not None):
                 self._connection = FakeConnection(client_info.user_agent)
 
-    @patch("google.cloud.automl_v1beta1.AutoMlClient", new=FakeClient)
+    @patch("google.cloud.automl.AutoMlClient", new=FakeClient)
     def test_user_provided_credentials(self):
         credentials = _make_credentials()
         env = EnvironmentVarGuard()
@@ -42,10 +37,10 @@ def test_user_provided_credentials(self):
     def test_tables_gcs_client(self):
         # The GcsClient can't currently be monkeypatched for default
         # credentials because it requires a project which can't be set.
-        # Verify that creating an automl.GcsClient given an actual
+        # Verify that creating an automl_v1beta1.GcsClient given an actual
         # storage.Client sets the client properly.
         gcs_client = storage.Client(project="xyz", credentials=_make_credentials())
-        tables_gcs_client = automl.GcsClient(client=gcs_client)
+        tables_gcs_client = automl_v1beta1.GcsClient(client=gcs_client)
         self.assertIs(tables_gcs_client.client, gcs_client)
 
     @patch("google.cloud.automl_v1beta1.gapic.auto_ml_client.AutoMlClient", new=FakeClient)
@@ -56,10 +51,10 @@ def test_tables_client_credentials(self):
         env.set('KAGGLE_KERNEL_INTEGRATIONS', 'AUTOML')
         with env:
             init_automl()
-            tables_client = automl.TablesClient(credentials=credentials)
+            tables_client = automl_v1beta1.TablesClient(credentials=credentials)
             self.assertEqual(tables_client.auto_ml_client.credentials, credentials)
 
-    @patch("google.cloud.automl_v1beta1.AutoMlClient", new=FakeClient)
+    @patch("google.cloud.automl.AutoMlClient", new=FakeClient)
     def test_default_credentials_automl_client(self):
         env = EnvironmentVarGuard()
         env.set('KAGGLE_USER_SECRETS_TOKEN', 'foobar')
@@ -71,18 +66,31 @@ def test_default_credentials_automl_client(self):
             self.assertIsInstance(automl_client.credentials, KaggleKernelCredentials)
             self.assertTrue(automl_client._connection.user_agent.startswith("kaggle-gcp-client/1.0"))
 
+    @patch("google.cloud.automl_v1beta1.AutoMlClient", new=FakeClient)
+    def test_default_credentials_automl_v1beta1_client(self):
+        env = EnvironmentVarGuard()
+        env.set('KAGGLE_USER_SECRETS_TOKEN', 'foobar')
+        env.set('KAGGLE_KERNEL_INTEGRATIONS', 'AUTOML')
+        with env:
+            init_automl()
+            automl_client = automl_v1beta1.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):
         env = EnvironmentVarGuard()
         env.set('KAGGLE_USER_SECRETS_TOKEN', 'foobar')
         env.set('KAGGLE_KERNEL_INTEGRATIONS', 'AUTOML')
         with env:
             init_automl()
-            tables_client = automl.TablesClient()
+            tables_client = automl_v1beta1.TablesClient()
             self.assertIsNotNone(tables_client.credentials)
             self.assertIsInstance(tables_client.credentials, KaggleKernelCredentials)
+            self.assertTrue(tables_client._connection.user_agent.startswith("kaggle-gcp-client/1.0"))
 
-    @patch("google.cloud.automl_v1beta1.PredictionServiceClient", new=FakeClient)
+    @patch("google.cloud.automl.PredictionServiceClient", new=FakeClient)
     def test_default_credentials_prediction_client(self):
         env = EnvironmentVarGuard()
         env.set('KAGGLE_USER_SECRETS_TOKEN', 'foobar')
@@ -91,6 +99,18 @@ def test_default_credentials_prediction_client(self):
             prediction_client = automl.PredictionServiceClient()
             self.assertIsNotNone(prediction_client.credentials)
             self.assertIsInstance(prediction_client.credentials, KaggleKernelCredentials)
+            self.assertTrue(prediction_client._connection.user_agent.startswith("kaggle-gcp-client/1.0"))
+
+    @patch("google.cloud.automl_v1beta1.PredictionServiceClient", new=FakeClient)
+    def test_default_credentials_prediction_v1beta1_client(self):
+        env = EnvironmentVarGuard()
+        env.set('KAGGLE_USER_SECRETS_TOKEN', 'foobar')
+        env.set('KAGGLE_KERNEL_INTEGRATIONS', 'AUTOML')
+        with env:
+            prediction_client = automl_v1beta1.PredictionServiceClient()
+            self.assertIsNotNone(prediction_client.credentials)
+            self.assertIsInstance(prediction_client.credentials, KaggleKernelCredentials)
+            self.assertTrue(prediction_client._connection.user_agent.startswith("kaggle-gcp-client/1.0"))
 
     def test_monkeypatching_idempotent(self):
         env = EnvironmentVarGuard()