Skip to content

Commit

Permalink
fix: use REST API in cell magic when requested (#892)
Browse files Browse the repository at this point in the history
Fixes #876.

The `--use_rest_api` option did not work as expected and this commit fixes it.

**PR checklist:**
- [x] Make sure to open an issue as a [bug/issue](https://github.com/googleapis/python-bigquery/issues/new/choose) before writing your code!  That way we can discuss the change, evaluate designs, and agree on the general idea
- [x] Ensure the tests and linter pass
- [x] Code coverage does not decrease (if any source code was changed)
- [x] Appropriate docs were updated (if necessary)
  • Loading branch information
plamut authored Aug 25, 2021
1 parent 72a52f0 commit 1cb3e55
Show file tree
Hide file tree
Showing 2 changed files with 25 additions and 8 deletions.
12 changes: 9 additions & 3 deletions google/cloud/bigquery/magics/magics.py
Original file line number Diff line number Diff line change
Expand Up @@ -671,7 +671,9 @@ def _cell_magic(line, query):
_handle_error(ex, args.destination_var)
return

result = rows.to_dataframe(bqstorage_client=bqstorage_client)
result = rows.to_dataframe(
bqstorage_client=bqstorage_client, create_bqstorage_client=False,
)
if args.destination_var:
IPython.get_ipython().push({args.destination_var: result})
return
Expand Down Expand Up @@ -728,11 +730,15 @@ def _cell_magic(line, query):

if max_results:
result = query_job.result(max_results=max_results).to_dataframe(
bqstorage_client=bqstorage_client, progress_bar_type=progress_bar
bqstorage_client=None,
create_bqstorage_client=False,
progress_bar_type=progress_bar,
)
else:
result = query_job.to_dataframe(
bqstorage_client=bqstorage_client, progress_bar_type=progress_bar
bqstorage_client=bqstorage_client,
create_bqstorage_client=False,
progress_bar_type=progress_bar,
)

if args.destination_var:
Expand Down
21 changes: 16 additions & 5 deletions tests/unit/test_magics.py
Original file line number Diff line number Diff line change
Expand Up @@ -660,7 +660,9 @@ def warning_match(warning):
assert client_info.user_agent == "ipython-" + IPython.__version__

query_job_mock.to_dataframe.assert_called_once_with(
bqstorage_client=bqstorage_instance_mock, progress_bar_type="tqdm"
bqstorage_client=bqstorage_instance_mock,
create_bqstorage_client=mock.ANY,
progress_bar_type="tqdm",
)

assert isinstance(return_value, pandas.DataFrame)
Expand Down Expand Up @@ -703,7 +705,9 @@ def test_bigquery_magic_with_rest_client_requested(monkeypatch):

bqstorage_mock.assert_not_called()
query_job_mock.to_dataframe.assert_called_once_with(
bqstorage_client=None, progress_bar_type="tqdm"
bqstorage_client=None,
create_bqstorage_client=False,
progress_bar_type="tqdm",
)

assert isinstance(return_value, pandas.DataFrame)
Expand Down Expand Up @@ -757,7 +761,12 @@ def test_bigquery_magic_w_max_results_valid_calls_queryjob_result():
client_query_mock.return_value = query_job_mock
ip.run_cell_magic("bigquery", "--max_results=5", sql)

query_job_mock.result.assert_called_with(max_results=5)
query_job_mock.result.assert_called_with(max_results=5)
query_job_mock.result.return_value.to_dataframe.assert_called_once_with(
bqstorage_client=None,
create_bqstorage_client=False,
progress_bar_type=mock.ANY,
)


@pytest.mark.usefixtures("ipython_interactive")
Expand Down Expand Up @@ -929,7 +938,7 @@ def test_bigquery_magic_w_table_id_and_bqstorage_client():

ip.run_cell_magic("bigquery", "--max_results=5", table_id)
row_iterator_mock.to_dataframe.assert_called_once_with(
bqstorage_client=bqstorage_instance_mock
bqstorage_client=bqstorage_instance_mock, create_bqstorage_client=mock.ANY,
)


Expand Down Expand Up @@ -1246,7 +1255,9 @@ def test_bigquery_magic_w_progress_bar_type_w_context_setter(monkeypatch):

bqstorage_mock.assert_not_called()
query_job_mock.to_dataframe.assert_called_once_with(
bqstorage_client=None, progress_bar_type=magics.context.progress_bar_type
bqstorage_client=None,
create_bqstorage_client=False,
progress_bar_type=magics.context.progress_bar_type,
)

assert isinstance(return_value, pandas.DataFrame)
Expand Down

0 comments on commit 1cb3e55

Please sign in to comment.