From bdef180caf9e222fef43e3eb7270b1c73d098e26 Mon Sep 17 00:00:00 2001 From: Peter Lamut Date: Mon, 23 Aug 2021 15:08:57 +0200 Subject: [PATCH 1/2] fix: use REST API in cell magic when requested The --use_rest_api option did not work as expected and this commit fixes it. --- google/cloud/bigquery/magics/magics.py | 13 ++++++++++--- tests/unit/test_magics.py | 21 ++++++++++++++++----- 2 files changed, 26 insertions(+), 8 deletions(-) diff --git a/google/cloud/bigquery/magics/magics.py b/google/cloud/bigquery/magics/magics.py index 2b8c2928e..849d8e97c 100644 --- a/google/cloud/bigquery/magics/magics.py +++ b/google/cloud/bigquery/magics/magics.py @@ -671,7 +671,10 @@ 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=use_bqstorage_api, + ) if args.destination_var: IPython.get_ipython().push({args.destination_var: result}) return @@ -728,11 +731,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=use_bqstorage_api, + progress_bar_type=progress_bar, ) if args.destination_var: diff --git a/tests/unit/test_magics.py b/tests/unit/test_magics.py index d030482cc..88c92a070 100644 --- a/tests/unit/test_magics.py +++ b/tests/unit/test_magics.py @@ -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) @@ -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) @@ -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") @@ -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, ) @@ -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) From a5b5b4cdd5230fed58aaa5eb38ce0a2394d7696d Mon Sep 17 00:00:00 2001 From: Peter Lamut Date: Tue, 24 Aug 2021 09:51:24 +0200 Subject: [PATCH 2/2] Always set create_bqstorage_client to False --- google/cloud/bigquery/magics/magics.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/google/cloud/bigquery/magics/magics.py b/google/cloud/bigquery/magics/magics.py index 849d8e97c..d368bbeaa 100644 --- a/google/cloud/bigquery/magics/magics.py +++ b/google/cloud/bigquery/magics/magics.py @@ -672,8 +672,7 @@ def _cell_magic(line, query): return result = rows.to_dataframe( - bqstorage_client=bqstorage_client, - create_bqstorage_client=use_bqstorage_api, + bqstorage_client=bqstorage_client, create_bqstorage_client=False, ) if args.destination_var: IPython.get_ipython().push({args.destination_var: result}) @@ -738,7 +737,7 @@ def _cell_magic(line, query): else: result = query_job.to_dataframe( bqstorage_client=bqstorage_client, - create_bqstorage_client=use_bqstorage_api, + create_bqstorage_client=False, progress_bar_type=progress_bar, )