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

Fixed subtle bug where map() function was lazy and never executed, so… · feibyte/docker-python@5d5d7a8 · GitHub

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

Commit 5d5d7a8

Browse files
vimota
authored
Fixed subtle bug where map() function was lazy and never executed, so we didn't pick up that the user had an integration connected. Also cleaned up the code to remove incorrect env var name and unused variables. (Kaggle#536)
1 parent 5c3ef7b commit 5d5d7a8

3 files changed

Lines changed: 45 additions & 32 deletions

File tree

‎patches/kaggle_gcp.py‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,18 +1,21 @@
11
import os
22
from google.auth import credentials
3-
from google.auth.exceptions import RefreshError
3+
from google.auth.exceptions import RefreshError
44
from google.cloud import bigquery
55
from google.cloud.bigquery._http import Connection
66
from kaggle_secrets import UserSecretsClient
77

8+
89
def get_integrations():
910
kernel_integrations_var = os.getenv("KAGGLE_KERNEL_INTEGRATIONS")
1011
kernel_integrations = KernelIntegrations()
1112
if kernel_integrations_var is None:
1213
return kernel_integrations
13-
map(lambda x: kernel_integrations.add_integration(x.lower()), kernel_integrations_var.split(':'))
14+
for integration in kernel_integrations_var.split(':'):
15+
kernel_integrations.add_integration(integration.lower())
1416
return kernel_integrations
1517

18+
1619
class KernelIntegrations():
1720
def __init__(self):
1821
self.integrations = {}

‎patches/sitecustomize.py‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@
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_BQ_USER_JWT")
6+
bq_user_jwt = os.getenv("KAGGLE_USER_SECRETS_TOKEN")
77
if kaggle_proxy_token or bq_user_jwt:
88
from google.auth import credentials
99
from google.cloud import bigquery
@@ -13,7 +13,6 @@
1313
from kaggle_gcp import PublicBigqueryClient, KaggleKernelCredentials
1414

1515
def monkeypatch_bq(bq_client, *args, **kwargs):
16-
data_proxy_project = os.getenv("KAGGLE_DATA_PROXY_PROJECT")
1716
specified_project = kwargs.get('project')
1817
specified_credentials = kwargs.get('credentials')
1918
kernel_integrations = get_integrations()

‎tests/test_bigquery.py‎

Lines changed: 39 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -12,106 +12,117 @@
1212

1313

1414
class TestBigQuery(unittest.TestCase):
15-
15+
1616
def _test_proxy(self, client, should_use_proxy):
1717
class HTTPHandler(BaseHTTPRequestHandler):
1818
called = False
1919
header_found = False
2020

21-
def do_HEAD(s):
22-
s.send_response(200)
21+
def do_HEAD(self):
22+
self.send_response(200)
2323

24-
def do_GET(s):
24+
def do_GET(self):
2525
HTTPHandler.called = True
26-
HTTPHandler.header_found = any(k for k in s.headers if k == "X-KAGGLE-PROXY-DATA" and s.headers[k] == "test-key")
27-
s.send_response(200)
26+
HTTPHandler.header_found = any(
27+
k for k in self.headers if k == "X-KAGGLE-PROXY-DATA" and self.headers[k] == "test-key")
28+
self.send_response(200)
2829

2930
server_address = urlparse(os.getenv('KAGGLE_DATA_PROXY_URL'))
3031
with HTTPServer((server_address.hostname, server_address.port), HTTPHandler) as httpd:
3132
threading.Thread(target=httpd.serve_forever).start()
3233

3334
try:
34-
for ds in client.list_datasets(): pass
35+
for _ in client.list_datasets():
36+
pass
3537
except:
3638
pass
3739

3840
httpd.shutdown()
3941
if should_use_proxy:
40-
self.assertTrue(HTTPHandler.called, msg="Fake server did not receive a request from the BQ client.")
41-
self.assertTrue(HTTPHandler.header_found, msg="X-KAGGLE-PROXY-DATA header was missing from the BQ request.")
42+
self.assertTrue(
43+
HTTPHandler.called, msg="Fake server did not receive a request from the BQ client.")
44+
self.assertTrue(
45+
HTTPHandler.header_found, msg="X-KAGGLE-PROXY-DATA header was missing from the BQ request.")
4246
else:
43-
self.assertFalse(HTTPHandler.called, msg="Fake server was called from the BQ client, but should not have been.")
47+
self.assertFalse(
48+
HTTPHandler.called, msg="Fake server was called from the BQ client, but should not have been.")
4449

4550
def test_proxy_using_library(self):
4651
env = EnvironmentVarGuard()
47-
env.unset('KAGGLE_BQ_USER_JWT')
52+
env.unset('KAGGLE_USER_SECRETS_TOKEN')
4853
with env:
4954
client = PublicBigqueryClient()
5055
self._test_proxy(client, should_use_proxy=True)
5156

5257
def test_proxy_no_project(self):
5358
env = EnvironmentVarGuard()
54-
env.unset('KAGGLE_BQ_USER_JWT')
59+
env.unset('KAGGLE_USER_SECRETS_TOKEN')
5560
with env:
5661
client = bigquery.Client()
5762
self._test_proxy(client, should_use_proxy=True)
5863

5964
def test_proxy_with_kwargs(self):
6065
env = EnvironmentVarGuard()
61-
env.unset('KAGGLE_BQ_USER_JWT')
66+
env.unset('KAGGLE_USER_SECRETS_TOKEN')
6267
with env:
63-
client = bigquery.Client(default_query_job_config=bigquery.QueryJobConfig(maximum_bytes_billed=1e9))
68+
client = bigquery.Client(
69+
default_query_job_config=bigquery.QueryJobConfig(maximum_bytes_billed=1e9))
6470
self._test_proxy(client, should_use_proxy=True)
6571

6672
def test_project_with_connected_account(self):
6773
env = EnvironmentVarGuard()
68-
env.set('KAGGLE_BQ_USER_JWT', 'foobar')
74+
env.set('KAGGLE_USER_SECRETS_TOKEN', 'foobar')
6975
with env:
70-
client = bigquery.Client(project='ANOTHER_PROJECT', credentials=KaggleKernelCredentials())
76+
client = bigquery.Client(
77+
project='ANOTHER_PROJECT', credentials=KaggleKernelCredentials())
7178
self._test_proxy(client, should_use_proxy=False)
7279

7380
def test_project_with_empty_integrations(self):
7481
env = EnvironmentVarGuard()
75-
env.set('KAGGLE_BQ_USER_JWT', 'foobar')
82+
env.set('KAGGLE_USER_SECRETS_TOKEN', 'foobar')
7683
env.set('KAGGLE_KERNEL_INTEGRATIONS', '')
7784
with env:
78-
client = bigquery.Client(project='ANOTHER_PROJECT', credentials=KaggleKernelCredentials())
85+
client = bigquery.Client(
86+
project='ANOTHER_PROJECT', credentials=KaggleKernelCredentials())
7987
self._test_proxy(client, should_use_proxy=False)
8088

81-
def test_project_with_connected_account_default_credentials(self):
89+
def test_project_with_connected_account_unrelated_integrations(self):
8290
env = EnvironmentVarGuard()
83-
env.set('KAGGLE_BQ_USER_JWT', 'foobar')
91+
env.set('KAGGLE_USER_SECRETS_TOKEN', 'foobar')
8492
env.set('KAGGLE_KERNEL_INTEGRATIONS', 'GCS:ANOTHER_ONE')
8593
with env:
86-
client = bigquery.Client(project='ANOTHER_PROJECT', credentials=KaggleKernelCredentials())
94+
client = bigquery.Client(
95+
project='ANOTHER_PROJECT', credentials=KaggleKernelCredentials())
8796
self._test_proxy(client, should_use_proxy=False)
8897

89-
def test_project_with_connected_account_unrelated_integrations(self):
98+
def test_project_with_connected_account_default_credentials(self):
9099
env = EnvironmentVarGuard()
91-
env.set('KAGGLE_BQ_USER_JWT', 'foobar')
100+
env.set('KAGGLE_USER_SECRETS_TOKEN', 'foobar')
92101
env.set('KAGGLE_KERNEL_INTEGRATIONS', 'BIGQUERY')
93102
with env:
94-
client = bigquery.Client(project='ANOTHER_PROJECT')
103+
client = bigquery.Client(project='ANOTHER_PROJECT')
95104
self._test_proxy(client, should_use_proxy=False)
96105

97106
def test_simultaneous_clients(self):
98107
env = EnvironmentVarGuard()
99-
env.set('KAGGLE_BQ_USER_JWT', 'foobar')
108+
env.set('KAGGLE_USER_SECRETS_TOKEN', 'foobar')
100109
with env:
101110
proxy_client = bigquery.Client()
102111
self._test_proxy(proxy_client, should_use_proxy=True)
103-
bq_client = bigquery.Client(project='ANOTHER_PROJECT', credentials=KaggleKernelCredentials())
112+
bq_client = bigquery.Client(
113+
project='ANOTHER_PROJECT', credentials=KaggleKernelCredentials())
104114
self._test_proxy(bq_client, should_use_proxy=False)
105115
# Verify that proxy client is still going to proxy to ensure global Connection
106116
# isn't being modified.
107117
self._test_proxy(proxy_client, should_use_proxy=True)
108118

109119
def test_no_project_with_connected_account(self):
110120
env = EnvironmentVarGuard()
111-
env.set('KAGGLE_BQ_USER_JWT', 'foobar')
121+
env.set('KAGGLE_USER_SECRETS_TOKEN', 'foobar')
122+
env.set('KAGGLE_KERNEL_INTEGRATIONS', 'BIGQUERY')
112123
with env:
113124
with self.assertRaises(DefaultCredentialsError):
114125
# TODO(vimota): Handle this case, either default to Kaggle Proxy or use some default project
115126
# by the user or throw a custom exception.
116-
client = bigquery.Client(credentials=KaggleKernelCredentials())
127+
client = bigquery.Client()
117128
self._test_proxy(client, should_use_proxy=False)

0 commit comments

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