Add --params option to %%bigquery magic

bigquery/google/cloud/bigquery/magics.py

@@ -39,6 +39,9 @@
         amount of time for the query to complete will not be cleared after the
         query is finished. By default, this information will be displayed but
         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.
     * ``<query>`` (required, cell argument):
         SQL query to run.
 
@@ -100,6 +103,7 @@
 
 from __future__ import print_function
 
+import ast
 import time
 from concurrent import futures
 
@@ -112,6 +116,7 @@
 
 import google.auth
 from google.cloud import bigquery
+from google.cloud.bigquery.dbapi._helpers import to_query_parameters
 
 
 class Context(object):
@@ -249,6 +254,13 @@ def _run_query(client, query, job_config=None):
           'amount of time for the query to finish. By default, this '
           'information will be displayed as the query runs, but will be '
           'cleared after the query is finished.'))
+@magic_arguments.argument(
+    '--params',
+    nargs='+',
+    default=None,
+    help=('Parameters to format the query string. If present, it should be a '
+          'parsable JSON string. The parsed dictionary will be used for string'
+          'replacement in the query'))
 def _cell_magic(line, query):
     """Underlying function for bigquery cell magic
 
@@ -265,9 +277,17 @@ def _cell_magic(line, query):
     """
     args = magic_arguments.parse_argstring(_cell_magic, line)
 
+    params = []
+    if args.params is not None:
+        try:
+            params = to_query_parameters(ast.literal_eval(''.join(args.params)))
+        except Exception:
+            raise SyntaxError('--params is not a correctly formatted JSON string')
+
     project = args.project or context.project
     client = bigquery.Client(project=project, credentials=context.credentials)
     job_config = bigquery.job.QueryJobConfig()
+    job_config.query_parameters = params
     job_config.use_legacy_sql = args.use_legacy_sql
     query_job = _run_query(client, query, job_config)
 

bigquery/tests/unit/test_magics.py

@@ -259,3 +259,33 @@ def test_bigquery_magic_with_project():
         assert client_used.project == 'specific-project'
         # context project should not change
         assert magics.context.project == 'general-project'
+
+
+@pytest.mark.usefixtures('ipython_interactive')
+@pytest.mark.skipif(pandas is None, reason='Requires `pandas`')
+def test_bigquery_magic_with_formatting_params():
+    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
+
+    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
+
+        ip.run_cell_magic('bigquery', '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 len(df) == len(result)    # verify row count
+    assert list(df) == list(result)  # verify column names
+