Add test for expanding a dictionary. Add usage examples

bigquery/google/cloud/bigquery/magics.py

@@ -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.
 
@@ -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
@@ -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')
 

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,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)
@@ -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