diff --git a/google/cloud/bigquery/ipython_magics/magics.py b/google/cloud/bigquery/ipython_magics/magics.py index 7128e32bf..ee4557416 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. - * ```` (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 dd4546d87..1a0579a24 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)