Skip to content

Commit bfd94af

Browse files
guillermo-carrascoguillermo-carrasco
committed
Add test for expanding a dictionary. Add usage examples
1 parent 0baa021 commit bfd94af

2 files changed

Lines changed: 51 additions & 9 deletions

File tree

‎bigquery/google/cloud/bigquery/magics.py‎

Lines changed: 14 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -41,7 +41,7 @@
4141
will be cleared after the query is finished.
4242
* ``--params <params dictionary>`` (optional, line argument):
4343
If present, the argument must be a parsable JSON string. This dictionary
44-
will be used to format values enclosed within {} in the query.
44+
will be used to format values preceeded by @ in the query.
4545
* ``<query>`` (required, cell argument):
4646
SQL query to run.
4747
@@ -99,6 +99,17 @@
9999
...: 1 Patricia 1568495
100100
...: 2 Elizabeth 1519946
101101
102+
In [5]: %%bigquery df --params {"num": 17}
103+
...: SELECT @num AS num
104+
Out[5]:
105+
...: num
106+
...: 0 17
107+
In [6]: params = {"num": 17}
108+
In [7]: %%bigquery df --params $params
109+
...: SELECT @num AS num
110+
Out[7]:
111+
...: num
112+
...: 0 17
102113
"""
103114

104115
from __future__ import print_function
@@ -116,7 +127,7 @@
116127

117128
import google.auth
118129
from google.cloud import bigquery
119-
from google.cloud.bigquery.dbapi._helpers import to_query_parameters
130+
from google.cloud.bigquery.dbapi import _helpers
120131

121132

122133
class Context(object):
@@ -280,7 +291,7 @@ def _cell_magic(line, query):
280291
params = []
281292
if args.params is not None:
282293
try:
283-
params = to_query_parameters(ast.literal_eval(''.join(args.params)))
294+
params = _helpers.to_query_parameters(ast.literal_eval(''.join(args.params)))
284295
except Exception:
285296
raise SyntaxError('--params is not a correctly formatted JSON string')
286297

‎bigquery/tests/unit/test_magics.py‎

Lines changed: 37 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -186,7 +186,7 @@ def test_bigquery_magic_with_result_saved_to_variable():
186186

187187
sql = 'SELECT 17 AS num'
188188
result = pandas.DataFrame([17], columns=['num'])
189-
assert 'myvariable' not in ip.user_ns
189+
assert 'df' not in ip.user_ns
190190

191191
run_query_patch = mock.patch(
192192
'google.cloud.bigquery.magics._run_query', autospec=True)
@@ -263,15 +263,15 @@ def test_bigquery_magic_with_project():
263263

264264
@pytest.mark.usefixtures('ipython_interactive')
265265
@pytest.mark.skipif(pandas is None, reason='Requires `pandas`')
266-
def test_bigquery_magic_with_formatting_params():
266+
def test_bigquery_magic_with_formatting_params_with_string():
267267
ip = IPython.get_ipython()
268268
ip.extension_manager.load_extension('google.cloud.bigquery')
269269
magics.context.credentials = mock.create_autospec(
270270
google.auth.credentials.Credentials, instance=True)
271271

272272
sql = 'SELECT @num AS num'
273273
result = pandas.DataFrame([17], columns=['num'])
274-
assert 'myvariable' not in ip.user_ns
274+
assert 'params_string_df' not in ip.user_ns
275275

276276
run_query_patch = mock.patch(
277277
'google.cloud.bigquery.magics._run_query', autospec=True)
@@ -281,11 +281,42 @@ def test_bigquery_magic_with_formatting_params():
281281
with run_query_patch as run_query_mock:
282282
run_query_mock.return_value = query_job_mock
283283

284-
ip.run_cell_magic('bigquery', 'df --params {"num":17}', sql)
284+
ip.run_cell_magic('bigquery', 'params_string_df --params {"num":17}', sql)
285285
run_query_mock.assert_called_once_with(mock.ANY, sql.format(num=17), mock.ANY)
286286

287-
assert 'df' in ip.user_ns # verify that variable exists
288-
df = ip.user_ns['df']
287+
assert 'params_string_df' in ip.user_ns # verify that variable exists
288+
df = ip.user_ns['params_string_df']
289289
assert len(df) == len(result) # verify row count
290290
assert list(df) == list(result) # verify column names
291291

292+
293+
@pytest.mark.usefixtures('ipython_interactive')
294+
@pytest.mark.skipif(pandas is None, reason='Requires `pandas`')
295+
def test_bigquery_magic_with_formatting_params_with_expanded_dict():
296+
ip = IPython.get_ipython()
297+
ip.extension_manager.load_extension('google.cloud.bigquery')
298+
magics.context.credentials = mock.create_autospec(
299+
google.auth.credentials.Credentials, instance=True)
300+
301+
sql = 'SELECT @num AS num'
302+
result = pandas.DataFrame([17], columns=['num'])
303+
assert 'params_dict_df' not in ip.user_ns
304+
305+
run_query_patch = mock.patch(
306+
'google.cloud.bigquery.magics._run_query', autospec=True)
307+
query_job_mock = mock.create_autospec(
308+
google.cloud.bigquery.job.QueryJob, instance=True)
309+
query_job_mock.to_dataframe.return_value = result
310+
with run_query_patch as run_query_mock:
311+
run_query_mock.return_value = query_job_mock
312+
313+
params = {"num": 17}
314+
# Insert dictionary into user namespace so that it can be expanded
315+
ip.user_ns['params'] = params
316+
ip.run_cell_magic('bigquery', 'params_dict_df --params $params', sql)
317+
run_query_mock.assert_called_once_with(mock.ANY, sql.format(num=17), mock.ANY)
318+
319+
assert 'params_dict_df' in ip.user_ns # verify that variable exists
320+
df = ip.user_ns['params_dict_df']
321+
assert len(df) == len(result) # verify row count
322+
assert list(df) == list(result) # verify column names

0 commit comments

Comments
 (0)