Skip to content

Commit

Permalink
Use a custom parser for cell magic arguments
Browse files Browse the repository at this point in the history
  • Loading branch information
plamut committed Aug 17, 2020
1 parent 07d7651 commit 4e30e4c
Show file tree
Hide file tree
Showing 2 changed files with 39 additions and 14 deletions.
50 changes: 37 additions & 13 deletions google/cloud/bigquery/ipython_magics/magics.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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__)
Expand Down Expand Up @@ -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(
Expand All @@ -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 "
Expand Down Expand Up @@ -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
Expand Down
3 changes: 2 additions & 1 deletion tests/unit/test_magics.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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)


Expand Down

0 comments on commit 4e30e4c

Please sign in to comment.