FazBrowse GitHub Viewer | Trending |
URL:
| Home
Tools: [Download Repo ZIP]   [Original HTTPS Page]

Add magics credentials monkeypatching if the kernel has bigquery inte… · feibyte/docker-python@b0cdd48 · GitHub

/ docker-python Public
forked from Kaggle/docker-python

Commit b0cdd48

Browse files
vimota
authored
Add magics credentials monkeypatching if the kernel has bigquery inte… (Kaggle#553)
1 parent 031dbe7 commit b0cdd48

2 files changed

Lines changed: 64 additions & 35 deletions

File tree

‎patches/sitecustomize.py‎

Lines changed: 44 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -3,41 +3,50 @@
33
kaggle_proxy_token = os.getenv("KAGGLE_DATA_PROXY_TOKEN")
44
kernel_integrations_var = os.getenv("KAGGLE_KERNEL_INTEGRATIONS")
55

6-
bq_user_jwt = os.getenv("KAGGLE_USER_SECRETS_TOKEN")
7-
if kaggle_proxy_token or bq_user_jwt:
8-
from google.auth import credentials, environment_vars
9-
from google.cloud import bigquery
10-
from google.cloud.bigquery._http import Connection
11-
# TODO: Update this to the correct kaggle.gcp path once we no longer inject modules
12-
# from the worker.
13-
from kaggle_gcp import PublicBigqueryClient, KaggleKernelCredentials
6+
def init():
7+
bq_user_jwt = os.getenv("KAGGLE_USER_SECRETS_TOKEN")
8+
if kaggle_proxy_token or bq_user_jwt:
9+
from google.auth import credentials, environment_vars
10+
from google.cloud import bigquery
11+
from google.cloud.bigquery._http import Connection
12+
# TODO: Update this to the correct kaggle.gcp path once we no longer inject modules
13+
# from the worker.
14+
from kaggle_gcp import PublicBigqueryClient, KaggleKernelCredentials
1415

15-
def monkeypatch_bq(bq_client, *args, **kwargs):
16-
specified_credentials = kwargs.get('credentials')
17-
has_bigquery = get_integrations().has_bigquery()
18-
# Prioritize passed in project id, but if it is missing look for env var.
19-
arg_project = kwargs.get('project')
20-
explicit_project_id = arg_project or os.environ.get(environment_vars.PROJECT)
21-
# This is a hack to get around the bug in google-cloud library.
22-
# Remove these two lines once this is resolved:
23-
# https://github.com/googleapis/google-cloud-python/issues/8108
24-
if explicit_project_id:
25-
kwargs['project'] = explicit_project_id
26-
if explicit_project_id is None and specified_credentials is None and not has_bigquery:
27-
print("Using Kaggle's public dataset BigQuery integration.")
28-
return PublicBigqueryClient(*args, **kwargs)
16+
# If this Kernel has bigquery integration on startup, preload the Kaggle Credentials
17+
# object for magics to work.
18+
if get_integrations().has_bigquery():
19+
from google.cloud.bigquery import magics
20+
magics.context.credentials = KaggleKernelCredentials()
2921

30-
else:
31-
if specified_credentials is None:
32-
kwargs['credentials'] = KaggleKernelCredentials()
33-
if (not has_bigquery):
34-
print('Please ensure you have selected a BigQuery '
35-
'account in the Kernels Settings sidebar.')
36-
return bq_client(*args, **kwargs)
22+
def monkeypatch_bq(bq_client, *args, **kwargs):
23+
specified_credentials = kwargs.get('credentials')
24+
has_bigquery = get_integrations().has_bigquery()
25+
# Prioritize passed in project id, but if it is missing look for env var.
26+
arg_project = kwargs.get('project')
27+
explicit_project_id = arg_project or os.environ.get(environment_vars.PROJECT)
28+
# This is a hack to get around the bug in google-cloud library.
29+
# Remove these two lines once this is resolved:
30+
# https://github.com/googleapis/google-cloud-python/issues/8108
31+
if explicit_project_id:
32+
kwargs['project'] = explicit_project_id
33+
if explicit_project_id is None and specified_credentials is None and not has_bigquery:
34+
print("Using Kaggle's public dataset BigQuery integration.")
35+
return PublicBigqueryClient(*args, **kwargs)
3736

38-
# Monkey patches BigQuery client creation to use proxy or user-connected GCP account.
39-
# Deprecated in favor of Kaggle.DataProxyClient().
40-
# TODO: Remove this once uses have migrated to that new interface.
41-
bq_client = bigquery.Client
42-
bigquery.Client = lambda *args, **kwargs: monkeypatch_bq(
43-
bq_client, *args, **kwargs)
37+
else:
38+
if specified_credentials is None:
39+
kwargs['credentials'] = KaggleKernelCredentials()
40+
if (not has_bigquery):
41+
print('Please ensure you have selected a BigQuery '
42+
'account in the Kernels Settings sidebar.')
43+
return bq_client(*args, **kwargs)
44+
45+
# Monkey patches BigQuery client creation to use proxy or user-connected GCP account.
46+
# Deprecated in favor of Kaggle.DataProxyClient().
47+
# TODO: Remove this once uses have migrated to that new interface.
48+
bq_client = bigquery.Client
49+
bigquery.Client = lambda *args, **kwargs: monkeypatch_bq(
50+
bq_client, *args, **kwargs)
51+
52+
init()

‎tests/test_bigquery.py‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -135,3 +135,23 @@ def test_no_project_with_connected_account(self):
135135
# by the user or throw a custom exception.
136136
client = bigquery.Client()
137137
self._test_proxy(client, should_use_proxy=False)
138+
139+
def test_magics_with_connected_account_default_credentials(self):
140+
env = EnvironmentVarGuard()
141+
env.set('KAGGLE_USER_SECRETS_TOKEN', 'foobar')
142+
env.set('KAGGLE_KERNEL_INTEGRATIONS', 'BIGQUERY')
143+
with env:
144+
import sitecustomize
145+
sitecustomize.init()
146+
from google.cloud.bigquery import magics
147+
self.assertEqual(type(magics.context._credentials), KaggleKernelCredentials)
148+
magics.context.credentials = None
149+
150+
def test_magics_without_connected_account(self):
151+
env = EnvironmentVarGuard()
152+
env.set('KAGGLE_USER_SECRETS_TOKEN', 'foobar')
153+
with env:
154+
import sitecustomize
155+
sitecustomize.init()
156+
from google.cloud.bigquery import magics
157+
self.assertIsNone(magics.context._credentials)

0 commit comments

Comments
 (0)
Back | FazBrowse Home | New Git URL