diff --git a/bigquery/google/cloud/bigquery/magics.py b/bigquery/google/cloud/bigquery/magics.py index 4494d7ba5ad1..27626752d8e7 100644 --- a/bigquery/google/cloud/bigquery/magics.py +++ b/bigquery/google/cloud/bigquery/magics.py @@ -32,6 +32,11 @@ * ``--project `` (optional, line argument): Project to use for running the query. Defaults to the context :attr:`~google.cloud.bigquery.magics.Context.project`. + * ``--use_bqstorage_api`` (optional, line argument): + Downloads the DataFrame using the BigQuery Storage API. To use this + option, install the ``google-cloud-bigquery-storage`` and ``fastavro`` + packages, and `enable the BigQuery Storage API + `_. * ``--use_legacy_sql`` (optional, line argument): Runs the query using Legacy SQL syntax. Defaults to Standard SQL if this argument not used. @@ -134,6 +139,10 @@ from IPython.core import magic_arguments except ImportError: # pragma: NO COVER raise ImportError("This module can only be loaded in IPython.") +try: + from google.cloud import bigquery_storage_v1beta1 +except ImportError: # pragma: NO COVER + bigquery_storage_v1beta1 = None import google.auth from google.cloud import bigquery @@ -150,6 +159,7 @@ class Context(object): def __init__(self): self._credentials = None self._project = None + self._use_bqstorage_api = None @property def credentials(self): @@ -211,6 +221,21 @@ def project(self): def project(self, value): self._project = value + @property + def use_bqstorage_api(self): + """bool: [Beta] Set to True to use the BigQuery Storage API to + download query results + + To use this option, install the ``google-cloud-bigquery-storage`` and + ``fastavro`` packages, and `enable the BigQuery Storage API + `_. + """ + return self._use_bqstorage_api + + @use_bqstorage_api.setter + def use_bqstorage_api(self, value): + self._use_bqstorage_api = value + context = Context() @@ -274,6 +299,16 @@ def _run_query(client, query, job_config=None): "Standard SQL if this argument is not used." ), ) +@magic_arguments.argument( + "--use_bqstorage_api", + action="store_true", + default=False, + help=( + "[Beta] Use the BigQuery Storage API to download large query results. " + "To use this option, install the google-cloud-bigquery-storage and " + "fastavro packages, and enable the BigQuery Storage API." + ), +) @magic_arguments.argument( "--verbose", action="store_true", @@ -328,6 +363,9 @@ def _cell_magic(line, query): project = args.project or context.project client = bigquery.Client(project=project, credentials=context.credentials) + bqstorage_client = _make_bqstorage_client( + args.use_bqstorage_api or context.use_bqstorage_api, context.credentials + ) job_config = bigquery.job.QueryJobConfig() job_config.query_parameters = params job_config.use_legacy_sql = args.use_legacy_sql @@ -336,8 +374,21 @@ def _cell_magic(line, query): if not args.verbose: display.clear_output() - result = query_job.to_dataframe() + result = query_job.to_dataframe(bqstorage_client=bqstorage_client) if args.destination_var: IPython.get_ipython().push({args.destination_var: result}) else: return result + + +def _make_bqstorage_client(use_bqstorage_api, credentials): + if not use_bqstorage_api: + return None + + if bigquery_storage_v1beta1 is None: + raise ImportError( + "Install the google-cloud-bigquery-storage and fastavro packages " + "to use the BigQuery Storage API." + ) + + return bigquery_storage_v1beta1.BigQueryStorageClient(credentials=credentials) diff --git a/bigquery/tests/unit/test_magics.py b/bigquery/tests/unit/test_magics.py index 6ec6123e726c..0db0bc5ebf52 100644 --- a/bigquery/tests/unit/test_magics.py +++ b/bigquery/tests/unit/test_magics.py @@ -31,6 +31,11 @@ IPython = None import google.auth.credentials + +try: + from google.cloud import bigquery_storage_v1beta1 +except ImportError: # pragma: NO COVER + bigquery_storage_v1beta1 = None from google.cloud.bigquery import table from google.cloud.bigquery import magics @@ -129,6 +134,37 @@ def test__run_query(): assert re.match("Query complete after .*s", updates[-1]) +def test__make_bqstorage_client_false(): + credentials_mock = mock.create_autospec( + google.auth.credentials.Credentials, instance=True + ) + got = magics._make_bqstorage_client(False, credentials_mock) + assert got is None + + +@pytest.mark.skipIf( + bigquery_storage_v1beta1 is None, "Requires `google-cloud-bigquery-storage`" +) +def test__make_bqstorage_client_true(): + credentials_mock = mock.create_autospec( + google.auth.credentials.Credentials, instance=True + ) + got = magics._make_bqstorage_client(True, credentials_mock) + assert isinstance(got, bigquery_storage_v1beta1.BigQueryStorageClient) + + +def test__make_bqstorage_client_true_raises_import_error(monkeypatch): + monkeypatch.setattr(magics, "bigquery_storage_v1beta1", None) + credentials_mock = mock.create_autospec( + google.auth.credentials.Credentials, instance=True + ) + + with pytest.raises(ImportError) as exc: + magics._make_bqstorage_client(True, credentials_mock) + + assert "google-cloud-bigquery-storage" in str(exc) + + @pytest.mark.usefixtures("ipython_interactive") def test_extension_load(): ip = IPython.get_ipython() @@ -141,13 +177,16 @@ def test_extension_load(): @pytest.mark.usefixtures("ipython_interactive") @pytest.mark.skipif(pandas is None, reason="Requires `pandas`") -def test_bigquery_magic_without_optional_arguments(): +def test_bigquery_magic_without_optional_arguments(monkeypatch): ip = IPython.get_ipython() ip.extension_manager.load_extension("google.cloud.bigquery") magics.context.credentials = mock.create_autospec( google.auth.credentials.Credentials, instance=True ) + # Shouldn't fail when BigQuery Storage client isn't installed. + monkeypatch.setattr(magics, "bigquery_storage_v1beta1", None) + sql = "SELECT 17 AS num" result = pandas.DataFrame([17], columns=["num"]) run_query_patch = mock.patch( @@ -257,6 +296,149 @@ def test_bigquery_magic_clears_display_in_verbose_mode(): assert clear_mock.call_count == 1 +@pytest.mark.usefixtures("ipython_interactive") +@pytest.mark.skipIf( + bigquery_storage_v1beta1 is None, "Requires `google-cloud-bigquery-storage`" +) +def test_bigquery_magic_with_bqstorage_from_argument(monkeypatch): + ip = IPython.get_ipython() + ip.extension_manager.load_extension("google.cloud.bigquery") + mock_credentials = mock.create_autospec( + google.auth.credentials.Credentials, instance=True + ) + + # Set up the context with monkeypatch so that it's reset for subsequent + # tests. + monkeypatch.setattr(magics.context, "credentials", mock_credentials) + monkeypatch.setattr(magics.context, "use_bqstorage_api", False) + + # Mock out the BigQuery Storage API. + bqstorage_mock = mock.create_autospec( + bigquery_storage_v1beta1.BigQueryStorageClient + ) + bqstorage_instance_mock = mock.create_autospec( + bigquery_storage_v1beta1.BigQueryStorageClient, instance=True + ) + bqstorage_mock.return_value = bqstorage_instance_mock + monkeypatch.setattr( + magics.bigquery_storage_v1beta1, "BigQueryStorageClient", bqstorage_mock + ) + + sql = "SELECT 17 AS num" + result = pandas.DataFrame([17], columns=["num"]) + run_query_patch = mock.patch( + "google.cloud.bigquery.magics._run_query", autospec=True + ) + query_job_mock = mock.create_autospec( + google.cloud.bigquery.job.QueryJob, instance=True + ) + query_job_mock.to_dataframe.return_value = result + with run_query_patch as run_query_mock: + run_query_mock.return_value = query_job_mock + + return_value = ip.run_cell_magic("bigquery", "--use_bqstorage_api", sql) + + bqstorage_mock.assert_called_once_with(credentials=mock_credentials) + query_job_mock.to_dataframe.assert_called_once_with( + bqstorage_client=bqstorage_instance_mock + ) + + assert isinstance(return_value, pandas.DataFrame) + + +@pytest.mark.usefixtures("ipython_interactive") +@pytest.mark.skipIf( + bigquery_storage_v1beta1 is None, "Requires `google-cloud-bigquery-storage`" +) +def test_bigquery_magic_with_bqstorage_from_context(monkeypatch): + ip = IPython.get_ipython() + ip.extension_manager.load_extension("google.cloud.bigquery") + mock_credentials = mock.create_autospec( + google.auth.credentials.Credentials, instance=True + ) + + # Set up the context with monkeypatch so that it's reset for subsequent + # tests. + monkeypatch.setattr(magics.context, "credentials", mock_credentials) + monkeypatch.setattr(magics.context, "use_bqstorage_api", True) + + # Mock out the BigQuery Storage API. + bqstorage_mock = mock.create_autospec( + bigquery_storage_v1beta1.BigQueryStorageClient + ) + bqstorage_instance_mock = mock.create_autospec( + bigquery_storage_v1beta1.BigQueryStorageClient, instance=True + ) + bqstorage_mock.return_value = bqstorage_instance_mock + monkeypatch.setattr( + magics.bigquery_storage_v1beta1, "BigQueryStorageClient", bqstorage_mock + ) + + sql = "SELECT 17 AS num" + result = pandas.DataFrame([17], columns=["num"]) + run_query_patch = mock.patch( + "google.cloud.bigquery.magics._run_query", autospec=True + ) + query_job_mock = mock.create_autospec( + google.cloud.bigquery.job.QueryJob, instance=True + ) + query_job_mock.to_dataframe.return_value = result + with run_query_patch as run_query_mock: + run_query_mock.return_value = query_job_mock + + return_value = ip.run_cell_magic("bigquery", "", sql) + + bqstorage_mock.assert_called_once_with(credentials=mock_credentials) + query_job_mock.to_dataframe.assert_called_once_with( + bqstorage_client=bqstorage_instance_mock + ) + + assert isinstance(return_value, pandas.DataFrame) + + +@pytest.mark.usefixtures("ipython_interactive") +@pytest.mark.skipIf( + bigquery_storage_v1beta1 is None, "Requires `google-cloud-bigquery-storage`" +) +def test_bigquery_magic_without_bqstorage(monkeypatch): + ip = IPython.get_ipython() + ip.extension_manager.load_extension("google.cloud.bigquery") + mock_credentials = mock.create_autospec( + google.auth.credentials.Credentials, instance=True + ) + + # Set up the context with monkeypatch so that it's reset for subsequent + # tests. + monkeypatch.setattr(magics.context, "credentials", mock_credentials) + + # Mock out the BigQuery Storage API. + bqstorage_mock = mock.create_autospec( + bigquery_storage_v1beta1.BigQueryStorageClient + ) + monkeypatch.setattr( + magics.bigquery_storage_v1beta1, "BigQueryStorageClient", bqstorage_mock + ) + + sql = "SELECT 17 AS num" + result = pandas.DataFrame([17], columns=["num"]) + run_query_patch = mock.patch( + "google.cloud.bigquery.magics._run_query", autospec=True + ) + query_job_mock = mock.create_autospec( + google.cloud.bigquery.job.QueryJob, instance=True + ) + query_job_mock.to_dataframe.return_value = result + with run_query_patch as run_query_mock: + run_query_mock.return_value = query_job_mock + + return_value = ip.run_cell_magic("bigquery", "", sql) + + bqstorage_mock.assert_not_called() + query_job_mock.to_dataframe.assert_called_once_with(bqstorage_client=None) + + assert isinstance(return_value, pandas.DataFrame) + + @pytest.mark.usefixtures("ipython_interactive") def test_bigquery_magic_with_project(): ip = IPython.get_ipython()