Separate Data Proxy monkeypatching into a distinct kaggle library and only expose gcp connected account as a credential class the user can pass into the vanilla bq client. In the future we should remove all monkeypatching and let users use either vanilla bq or the Kaggle library.

Author: vimota

Dockerfile

@@ -504,6 +504,7 @@ RUN pip install --upgrade dask && \
 
 # Add BigQuery client proxy settings
 ENV PYTHONUSERBASE "/root/.local"
+ADD patches/kaggle.py /root/.local/lib/python3.6/site-packages/kaggle.py
 ADD patches/sitecustomize.py /root/.local/lib/python3.6/site-packages/sitecustomize.py
 
 # Set backend for matplotlib

patches/kaggle.py

@@ -0,0 +1,40 @@
+import os
+from google.auth import credentials
+from google.cloud import bigquery
+from google.cloud.bigquery._http import Connection
+
+
+class KaggleKernelCredentials(credentials.Credentials):
+        """Custom Credentials used to authenticate using the Kernel's connected OAuth account."""
+
+        def refresh(self, request):
+            print("Calling Kaggle.UserSecrets to refresh token.")
+            # Set self.token and self.expiry here.
+            raise NotImplementedError(
+                "Private BigQuery integration is not yet implemented.")
+
+
+class DataProxyConnection(Connection):
+        """Custom Connection class used to proxy the BigQuery client ho Kaggle's data proxy."""
+
+        API_BASE_URL = os.getenv("KAGGLE_DATA_PROXY_URL")
+
+        def __init__(self, client):
+            super().__init__(client)
+            self._EXTRA_HEADERS["X-KAGGLE-PROXY-DATA"] = os.getenv("KAGGLE_DATA_PROXY_TOKEN")
+
+
+class DataProxyClient(bigquery.client.Client):
+    def __init__(self, project=None):
+        if project:
+            raise Exception("In order to query a private BigQuery project, please connect a GCP account. "
+                            "Otherwise do not specify a project to use Kaggle's public dataset BigQuery integration.")
+        data_proxy_project = os.getenv("KAGGLE_DATA_PROXY_PROJECT")
+        anon_credentials = credentials.AnonymousCredentials()
+        anon_credentials.refresh = lambda *args: None
+        super().__init__(
+            project=data_proxy_project, credentials=anon_credentials
+        )
+        self._connection = DataProxyConnection(self)
+
+kaggle_bq_client = lambda *args, **kwargs:  DataProxyClient(*args, **kwargs)

patches/sitecustomize.py

@@ -1,58 +1,26 @@
-from google.auth import credentials
-from google.cloud import bigquery
-from google.cloud.bigquery._http import Connection
 import os
 
-
-class KaggleKernelCredentials(credentials.Credentials):
-    """Custom Credentials used to authenticate using the Kernel's connected OAuth account."""
-
-    def refresh(self, request):
-        print("Calling Kaggle.UserSecrets to refresh token.")
-        # Set self.token and self.expiry here.
-        raise NotImplementedError("Private BigQuery integration is not yet implemented.")
-
 kaggle_proxy_token = os.getenv("KAGGLE_DATA_PROXY_TOKEN")
-CONNECTION_BASE_URL = Connection.API_BASE_URL
-
-
-def monkeypatch_bq(bq_client, *args, **kwargs):
-    data_proxy_project = os.getenv("KAGGLE_DATA_PROXY_PROJECT")
-    bq_user_jwt = os.getenv("KAGGLE_BQ_USER_JWT")
-    specified_project = kwargs.get('project')
-    # Use Data Proxy if user has specified to use the Kaggle project, or if
-    # there are no connected GCP accounts (to maintain backwards compatibility).
-    if bq_user_jwt is None and specified_project and specified_project.lower() != 'kaggle':
-        raise Exception("In order to query a private BigQuery project, please connect a GCP account. "
-                        "Otherwise specify 'kaggle' as the project to use Kaggle's public dataset BigQuery integration.")
-    use_data_proxy = (specified_project and specified_project.lower() == 'kaggle') or bq_user_jwt is None
-    if use_data_proxy:
-        if data_proxy_project is None or kaggle_proxy_token is None:
-            # We don't have the data proxy info so leave the bq client unmodified.
-            return bq_client(*args, **kwargs)
-        print("Using Kaggle's public dataset BigQuery integration.")
-        Connection.API_BASE_URL = os.getenv("KAGGLE_DATA_PROXY_URL")
-        Connection._EXTRA_HEADERS["X-KAGGLE-PROXY-DATA"] = kaggle_proxy_token
-        anon_credentials = credentials.AnonymousCredentials()
-        anon_credentials.refresh = lambda *args: None
-        kwargs['project'] = data_proxy_project
-        return bq_client(
-            *args,
-            credentials=anon_credentials,
-            **kwargs)
-    else:
-        Connection.API_BASE_URL = CONNECTION_BASE_URL
-        Connection._EXTRA_HEADERS.pop('X-KAGGLE-PROXY-DATA', None)
-        if kwargs.get('credentials') is not None:
-            # The user wants to use their own credentials scheme, don't try to interfere.
+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
+    from kaggle import kaggle_bq_client
+
+    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:
+            print("Using Kaggle's public dataset BigQuery integration.")
+            return kaggle_bq_client(*args, **kwargs)
+        else:
             return bq_client(*args, **kwargs)
-        print("Using enabled BigQuery integration.")
-        kwargs['credentials'] = KaggleKernelCredentials()
-        return bq_client(
-            *args,
-            **kwargs)
-
-# Monkey patches BigQuery client creation to use proxy or user.
-bq_client = bigquery.Client
-bigquery.Client = lambda *args, **kwargs:  monkeypatch_bq(bq_client, *args, **kwargs)
 
+    # Monkey patches BigQuery client creation to use proxy or user-connected GCP account.
+    # TODO: Remove monkeypatching altogether and move to using a Kaggle library for Data Proxy and
+    # KaggleKernelCredentials with vanilla bq client for connected account.
+    bq_client = bigquery.Client
+    bigquery.Client = lambda *args, **kwargs:  monkeypatch_bq(
+        bq_client, *args, **kwargs)

tests/test_bigquery.py

@@ -7,6 +7,7 @@
 
 from google.cloud import bigquery
 from google.auth.exceptions import DefaultCredentialsError
+from kaggle import KaggleKernelCredentials, kaggle_bq_client
 
 HOSTNAME = "127.0.0.1"
 PORT = 8000
@@ -41,12 +42,12 @@ def do_GET(s):
                 self.assertTrue(HTTPHandler.header_found, msg="X-KAGGLE-PROXY-DATA header was missing from the BQ request.")
             else:
                 self.assertFalse(HTTPHandler.called, msg="Fake server was called from the BQ client, but should not have been.")
-    
-    def test_proxy_kaggle_project(self):
+
+    def test_proxy_using_library(self):
         env = EnvironmentVarGuard()
         env.unset('KAGGLE_BQ_USER_JWT')
         with env:
-            client = bigquery.Client(project='KAGGLE')
+            client = kaggle_bq_client()
             self._test_proxy(client, should_use_proxy=True)
 
     def test_proxy_no_project(self):
@@ -60,8 +61,20 @@ def test_project_with_connected_account(self):
         env = EnvironmentVarGuard()
         env.set('KAGGLE_BQ_USER_JWT', 'foobar')
         with env:
-            client = bigquery.Client(project='ANOTHER_PROJECT')
+            client = bigquery.Client(project='ANOTHER_PROJECT', credentials=KaggleKernelCredentials())
             self._test_proxy(client, should_use_proxy=False)
+    
+    def test_simultaneous_clients(self):
+        env = EnvironmentVarGuard()
+        env.set('KAGGLE_BQ_USER_JWT', 'foobar')
+        with env:
+            proxy_client = bigquery.Client()
+            self._test_proxy(proxy_client, should_use_proxy=True)
+            bq_client = bigquery.Client(project='ANOTHER_PROJECT', credentials=KaggleKernelCredentials())
+            self._test_proxy(bq_client, should_use_proxy=False)
+            # Verify that proxy client is still going to proxy to ensure global Connection
+            # isn't being modified.
+            self._test_proxy(proxy_client, should_use_proxy=True)
 
     def test_no_project_with_connected_account(self):
         env = EnvironmentVarGuard()
@@ -70,5 +83,5 @@ def test_no_project_with_connected_account(self):
             with self.assertRaises(DefaultCredentialsError):
                 # TODO(vimota): Handle this case, either default to Kaggle Proxy or use some default project
                 # by the user or throw a custom exception.
-                client = bigquery.Client()
+                client = bigquery.Client(credentials=KaggleKernelCredentials())
                 self._test_proxy(client, should_use_proxy=False)