Skip to content

Commit

Permalink
Add test for expanding a dictionary. Add usage examples
Browse files Browse the repository at this point in the history
  • Loading branch information
guillermo-carrasco committed Oct 29, 2018
1 parent 0baa021 commit 2c17f19
Show file tree
Hide file tree
Showing 2 changed files with 52 additions and 9 deletions.
18 changes: 15 additions & 3 deletions bigquery/google/cloud/bigquery/magics.py
Original file line number Diff line number Diff line change
Expand Up @@ -41,7 +41,7 @@
will be cleared after the query is finished.
* ``--params <params dictionary>`` (optional, line argument):
If present, the argument must be a parsable JSON string. This dictionary
will be used to format values enclosed within {} in the query.
will be used to format values preceded by ``@`` in the query.
* ``<query>`` (required, cell argument):
SQL query to run.
Expand Down Expand Up @@ -99,6 +99,18 @@
...: 1 Patricia 1568495
...: 2 Elizabeth 1519946
In [5]: %%bigquery df --params {"num": 17}
...: SELECT @num AS num
Out[5]:
...: num
...: 0 17
In [6]: # Expand a dictionary instead of writing it's string value
In [6]: params = {"num": 17}
In [7]: %%bigquery df --params $params
...: SELECT @num AS num
Out[7]:
...: num
...: 0 17
"""

from __future__ import print_function
Expand All @@ -116,7 +128,7 @@

import google.auth
from google.cloud import bigquery
from google.cloud.bigquery.dbapi._helpers import to_query_parameters
from google.cloud.bigquery.dbapi import _helpers


class Context(object):
Expand Down Expand Up @@ -280,7 +292,7 @@ def _cell_magic(line, query):
params = []
if args.params is not None:
try:
params = to_query_parameters(ast.literal_eval(''.join(args.params)))
params = _helpers.to_query_parameters(ast.literal_eval(''.join(args.params)))
except Exception:
raise SyntaxError('--params is not a correctly formatted JSON string')

Expand Down
43 changes: 37 additions & 6 deletions bigquery/tests/unit/test_magics.py
Original file line number Diff line number Diff line change
Expand Up @@ -186,7 +186,7 @@ def test_bigquery_magic_with_result_saved_to_variable():

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

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

@pytest.mark.usefixtures('ipython_interactive')
@pytest.mark.skipif(pandas is None, reason='Requires `pandas`')
def test_bigquery_magic_with_formatting_params():
def test_bigquery_magic_with_formatting_params_with_string():
ip = IPython.get_ipython()
ip.extension_manager.load_extension('google.cloud.bigquery')
magics.context.credentials = mock.create_autospec(
google.auth.credentials.Credentials, instance=True)

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

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

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

assert 'df' in ip.user_ns # verify that variable exists
df = ip.user_ns['df']
assert 'params_string_df' in ip.user_ns # verify that variable exists
df = ip.user_ns['params_string_df']
assert len(df) == len(result) # verify row count
assert list(df) == list(result) # verify column names


@pytest.mark.usefixtures('ipython_interactive')
@pytest.mark.skipif(pandas is None, reason='Requires `pandas`')
def test_bigquery_magic_with_formatting_params_with_expanded_dict():
ip = IPython.get_ipython()
ip.extension_manager.load_extension('google.cloud.bigquery')
magics.context.credentials = mock.create_autospec(
google.auth.credentials.Credentials, instance=True)

sql = 'SELECT @num AS num'
result = pandas.DataFrame([17], columns=['num'])
assert 'params_dict_df' not in ip.user_ns

run_query_patch = mock.patch(
'google.cloud.bigquery.magics._run_query', autospec=True)
query_job_mock = mock.create_autospec(
google.cloud.bigquery.job.QueryJob, instance=True)
query_job_mock.to_dataframe.return_value = result
with run_query_patch as run_query_mock:
run_query_mock.return_value = query_job_mock

params = {"num": 17}
# Insert dictionary into user namespace so that it can be expanded
ip.user_ns['params'] = params
ip.run_cell_magic('bigquery', 'params_dict_df --params $params', sql)
run_query_mock.assert_called_once_with(mock.ANY, sql.format(num=17), mock.ANY)

assert 'params_dict_df' in ip.user_ns # verify that variable exists
df = ip.user_ns['params_dict_df']
assert len(df) == len(result) # verify row count
assert list(df) == list(result) # verify column names

0 comments on commit 2c17f19

Please sign in to comment.