From 4e30e4c552586fa23afbd3dcddf1acedd3e0fdff Mon Sep 17 00:00:00 2001
From: Peter Lamut <inbox@peterlamut.com>
Date: Mon, 17 Aug 2020 13:31:20 +0200
Subject: [PATCH] Use a custom parser for cell magic arguments

---
 .../cloud/bigquery/ipython_magics/magics.py   | 50 ++++++++++++++-----
 tests/unit/test_magics.py                     |  3 +-
 2 files changed, 39 insertions(+), 14 deletions(-)

diff --git a/google/cloud/bigquery/ipython_magics/magics.py b/google/cloud/bigquery/ipython_magics/magics.py
index 7128e32bfe..ee4557416d 100644
--- a/google/cloud/bigquery/ipython_magics/magics.py
+++ b/google/cloud/bigquery/ipython_magics/magics.py
@@ -65,13 +65,6 @@
           the variable name (ex. ``$my_dict_var``). See ``In[6]`` and ``In[7]``
           in the Examples section below.
 
-        .. note::
-
-            Due to the way IPython argument parser works, negative numbers in
-            dictionaries are incorrectly "recognized" as additional arguments,
-            resulting in an error ("unrecognized arguments"). To get around this,
-            pass such dictionary as a JSON string variable.
-
     * ``<query>`` (required, cell argument):
         SQL query to run. If the query does not contain any whitespace (aside
         from leading and trailing whitespace), it is assumed to represent a
@@ -159,13 +152,15 @@
 except ImportError:  # pragma: NO COVER
     raise ImportError("This module can only be loaded in IPython.")
 
+import six
+
 from google.api_core import client_info
 from google.api_core.exceptions import NotFound
 import google.auth
 from google.cloud import bigquery
 import google.cloud.bigquery.dataset
 from google.cloud.bigquery.dbapi import _helpers
-import six
+from google.cloud.bigquery.ipython_magics import line_arg_parser as lap
 
 
 IPYTHON_USER_AGENT = "ipython-{}".format(IPython.__version__)
@@ -473,7 +468,11 @@ def _cell_magic(line, query):
     Returns:
         pandas.DataFrame: the query results.
     """
-    args = magic_arguments.parse_argstring(_cell_magic, line)
+    # The built-in parser does not recognize Python structures such as dicts, thus
+    # we extract the "--params" option and inteprpret it separately.
+    params_option_value, rest_of_args = _split_args_line(line)
+
+    args = magic_arguments.parse_argstring(_cell_magic, rest_of_args)
 
     if args.use_bqstorage_api is not None:
         warnings.warn(
@@ -484,11 +483,17 @@ def _cell_magic(line, query):
     use_bqstorage_api = not args.use_rest_api
 
     params = []
-    if args.params is not None:
-        try:
-            params = _helpers.to_query_parameters(
-                ast.literal_eval("".join(args.params))
+    if params_option_value:
+        # A non-existing params variable is not expanded and ends up in the input
+        # in its raw form, e.g. "$query_params".
+        if params_option_value.startswith("$"):
+            msg = 'Parameter expansion failed, undefined variable "{}".'.format(
+                params_option_value[1:]
             )
+            raise NameError(msg)
+
+        try:
+            params = _helpers.to_query_parameters(ast.literal_eval(params_option_value))
         except Exception:
             raise SyntaxError(
                 "--params is not a correctly formatted JSON string or a JSON "
@@ -598,6 +603,25 @@ def _cell_magic(line, query):
         close_transports()
 
 
+def _split_args_line(line):
+    """Split out the --params option value from the input line arguments.
+
+    Args:
+        line (str): The line arguments passed to the cell magic.
+
+    Returns:
+        Tuple[str, str]
+    """
+    lexer = lap.Lexer(line)
+    scanner = lap.Parser(lexer)
+    tree = scanner.input_line()
+
+    extractor = lap.QueryParamsExtractor()
+    params_option_value, rest_of_args = extractor.visit(tree)
+
+    return params_option_value, rest_of_args
+
+
 def _make_bqstorage_client(use_bqstorage_api, credentials):
     if not use_bqstorage_api:
         return None
diff --git a/tests/unit/test_magics.py b/tests/unit/test_magics.py
index dd4546d87f..1a0579a24d 100644
--- a/tests/unit/test_magics.py
+++ b/tests/unit/test_magics.py
@@ -43,6 +43,7 @@
 from google.cloud import bigquery
 from google.cloud.bigquery import job
 from google.cloud.bigquery import table
+from google.cloud.bigquery.ipython_magics import line_arg_parser as lap
 from google.cloud.bigquery.ipython_magics import magics
 from tests.unit.helpers import make_connection
 from test_utils.imports import maybe_fail_import
@@ -1244,7 +1245,7 @@ def test_bigquery_magic_with_improperly_formatted_params():
 
     sql = "SELECT @num AS num"
 
-    with pytest.raises(SyntaxError):
+    with pytest.raises(lap.exceptions.QueryParamsParseError):
         ip.run_cell_magic("bigquery", "--params {17}", sql)