-
Notifications
You must be signed in to change notification settings - Fork 1.5k
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
BigQuery: Add ability to pass in a table ID instead of a query to the %%bigquery magic. #9170
Changes from 1 commit
77e62c5
7a13ca4
d663104
abcd29c
6f2dd64
428368f
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -129,6 +129,7 @@ | |
|
||
from __future__ import print_function | ||
|
||
import re | ||
import ast | ||
import sys | ||
import time | ||
|
@@ -266,6 +267,15 @@ def default_query_job_config(self, value): | |
context = Context() | ||
|
||
|
||
def _print_error(error, destination_var=None): | ||
if destination_var: | ||
print( | ||
"Could not save output to variable '{}'.".format(destination_var), | ||
file=sys.stderr, | ||
) | ||
print("\nERROR:\n", error, file=sys.stderr) | ||
|
||
|
||
def _run_query(client, query, job_config=None): | ||
"""Runs a query while printing status updates | ||
|
||
|
@@ -434,6 +444,26 @@ def _cell_magic(line, query): | |
else: | ||
max_results = None | ||
|
||
error = None | ||
|
||
if not re.search(r"\s", query.rstrip()): | ||
table_id = query.rstrip() | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. BTW, stripping whitespace from the same value again is not necessary, we could store the stripped string into a variable the first time we do it. There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. fixed in d663104 |
||
|
||
try: | ||
rows = client.list_rows(table_id, max_results=max_results) | ||
except Exception as ex: | ||
error = str(ex) | ||
if error: | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. could this be simplified to return from the exception? There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. it definitely can. fixed in 7a13ca4 |
||
_print_error(error, args.destination_var) | ||
return | ||
|
||
result = rows.to_dataframe(bqstorage_client=bqstorage_client) | ||
if args.destination_var: | ||
IPython.get_ipython().push({args.destination_var: result}) | ||
return | ||
else: | ||
return result | ||
|
||
job_config = bigquery.job.QueryJobConfig() | ||
job_config.query_parameters = params | ||
job_config.use_legacy_sql = args.use_legacy_sql | ||
|
@@ -445,7 +475,6 @@ def _cell_magic(line, query): | |
value = int(args.maximum_bytes_billed) | ||
job_config.maximum_bytes_billed = value | ||
|
||
error = None | ||
try: | ||
query_job = _run_query(client, query, job_config=job_config) | ||
except Exception as ex: | ||
|
@@ -455,12 +484,7 @@ def _cell_magic(line, query): | |
display.clear_output() | ||
|
||
if error: | ||
if args.destination_var: | ||
print( | ||
"Could not save output to variable '{}'.".format(args.destination_var), | ||
file=sys.stderr, | ||
) | ||
print("\nERROR:\n", error, file=sys.stderr) | ||
_print_error(error, args.destination_var) | ||
return | ||
|
||
if args.dry_run and args.destination_var: | ||
|
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
What scenario is this protecting against? Is there a scenario where unicode strings will not remove space characters via rstrip?
What do we expect in the else case?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Related: Does this repository enforce code coverage? Is there a case where we test there being unstoppable whitespace and not taking this branch?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This condition is actually testing for queries that contain whitespace characters that aren't removed by rstrip (so any whitespace that isn't a trailing newline). The assumption being made (as described in #9105 ) is that anything containing whitespace is a SQL query and won't take this branch, while anything string without whitespace is a table_id and will take this branch.
So anything regular SQL query would fall into the else case. There are tests that check whether query strings without spaces will be interpreted as table IDs and a test that checks if a string without whitespace that isn't a valid table_id raises an appropriate error message.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I can add a comment in the code if that would be helpful.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I would also write this down in a comment, i.e. that anything without whitespace is assumed to be table identifier which triggers a different use case of the magic command.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
added comment in d663104