From 6a1ebe3e8bdab912d854030e3cda8f0d91a7a2a5 Mon Sep 17 00:00:00 2001 From: Guillermo Carrasco Date: Mon, 29 Oct 2018 15:39:44 +0100 Subject: [PATCH] Add test for expanding a dictionary. Add usage examples --- bigquery/google/cloud/bigquery/magics.py | 18 ++++++++-- bigquery/tests/unit/test_magics.py | 43 ++++++++++++++++++++---- 2 files changed, 52 insertions(+), 9 deletions(-) diff --git a/bigquery/google/cloud/bigquery/magics.py b/bigquery/google/cloud/bigquery/magics.py index 7756eaf7751ab..e3e751fe5debf 100644 --- a/bigquery/google/cloud/bigquery/magics.py +++ b/bigquery/google/cloud/bigquery/magics.py @@ -41,7 +41,7 @@ will be cleared after the query is finished. * ``--params `` (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 preceeded by @ in the query. * ```` (required, cell argument): SQL query to run. @@ -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 [5]: # Expand a dictionary instead of writing it's string value + In [5]: params = {"num": 17} + In [6]: %%bigquery df --params $params + ...: SELECT @num AS num + Out[6]: + ...: num + ...: 0 17 """ from __future__ import print_function @@ -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): @@ -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') diff --git a/bigquery/tests/unit/test_magics.py b/bigquery/tests/unit/test_magics.py index 25fc1feffb736..daaf002270a12 100644 --- a/bigquery/tests/unit/test_magics.py +++ b/bigquery/tests/unit/test_magics.py @@ -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) @@ -263,7 +263,7 @@ 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( @@ -271,7 +271,7 @@ def test_bigquery_magic_with_formatting_params(): 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) @@ -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