Uses KAGGLE_KERNEL_INTEGRATIONS to determine the best default credent… (Kaggle#518)

* Uses KAGGLE_KERNEL_INTEGRATIONS to determine the best default credentials/client to use for bigquery.

* Add more test cases for potential KERNEL_INTEGRATIONS values.

Author: vimota

patches/kaggle_gcp.py

@@ -5,6 +5,24 @@
 from google.cloud.bigquery._http import Connection
 from kaggle_secrets import UserSecretsClient
 
+def get_integrations():
+    kernel_integrations_var = os.getenv("KAGGLE_KERNEL_INTEGRATIONS")
+    kernel_integrations = KernelIntegrations()
+    if kernel_integrations_var is None:
+        return kernel_integrations
+    map(lambda x: kernel_integrations.add_integration(x.lower()), kernel_integrations_var.split(':')) 
+    return kernel_integrations
+
+class KernelIntegrations():
+    def __init__(self):
+        self.integrations = {}
+
+    def add_integration(self, integration_name):
+        self.integrations[integration_name] = True
+
+    def has_bigquery(self):
+        return 'bigquery' in self.integrations.keys()
+
 
 class KaggleKernelCredentials(credentials.Credentials):
     """Custom Credentials used to authenticate using the Kernel's connected OAuth account.

patches/sitecustomize.py

@@ -1,23 +1,29 @@
 import os
-
+from kaggle_gcp import get_integrations
 kaggle_proxy_token = os.getenv("KAGGLE_DATA_PROXY_TOKEN")
+kernel_integrations_var = os.getenv("KAGGLE_KERNEL_INTEGRATIONS")
+
 bq_user_jwt = os.getenv("KAGGLE_BQ_USER_JWT")
 if kaggle_proxy_token or bq_user_jwt:
     from google.auth import credentials
     from google.cloud import bigquery
     from google.cloud.bigquery._http import Connection
     # TODO: Update this to the correct kaggle.gcp path once we no longer inject modules
     # from the worker.
-    from kaggle_gcp import PublicBigqueryClient
+    from kaggle_gcp import PublicBigqueryClient, KaggleKernelCredentials
 
     def monkeypatch_bq(bq_client, *args, **kwargs):
         data_proxy_project = os.getenv("KAGGLE_DATA_PROXY_PROJECT")
         specified_project = kwargs.get('project')
         specified_credentials = kwargs.get('credentials')
-        if specified_project is None and specified_credentials is None:
+        kernel_integrations = get_integrations()
+        if specified_project is None and specified_credentials is None and not kernel_integrations.has_bigquery():
             print("Using Kaggle's public dataset BigQuery integration.")
             return PublicBigqueryClient(*args, **kwargs)
+
         else:
+            if specified_credentials is None:
+                kwargs['credentials'] = KaggleKernelCredentials()
             return bq_client(*args, **kwargs)
 
     # Monkey patches BigQuery client creation to use proxy or user-connected GCP account.

tests/test_bigquery.py

@@ -62,7 +62,31 @@ def test_project_with_connected_account(self):
         with env:
             client = bigquery.Client(project='ANOTHER_PROJECT', credentials=KaggleKernelCredentials())
             self._test_proxy(client, should_use_proxy=False)
-    
+
+    def test_project_with_empty_integrations(self):
+        env = EnvironmentVarGuard()
+        env.set('KAGGLE_BQ_USER_JWT', 'foobar')
+        env.set('KAGGLE_KERNEL_INTEGRATIONS', '')
+        with env:
+            client = bigquery.Client(project='ANOTHER_PROJECT', credentials=KaggleKernelCredentials())
+            self._test_proxy(client, should_use_proxy=False)
+
+    def test_project_with_connected_account_default_credentials(self):
+        env = EnvironmentVarGuard()
+        env.set('KAGGLE_BQ_USER_JWT', 'foobar')
+        env.set('KAGGLE_KERNEL_INTEGRATIONS', 'GCS:ANOTHER_ONE')
+        with env:
+            client = bigquery.Client(project='ANOTHER_PROJECT', credentials=KaggleKernelCredentials())
+            self._test_proxy(client, should_use_proxy=False)
+
+    def test_project_with_connected_account_unrelated_integrations(self):
+        env = EnvironmentVarGuard()
+        env.set('KAGGLE_BQ_USER_JWT', 'foobar')
+        env.set('KAGGLE_KERNEL_INTEGRATIONS', 'BIGQUERY')
+        with env:
+            client = bigquery.Client(project='ANOTHER_PROJECT') 
+            self._test_proxy(client, should_use_proxy=False)
+
     def test_simultaneous_clients(self):
         env = EnvironmentVarGuard()
         env.set('KAGGLE_BQ_USER_JWT', 'foobar')