tswast · Apr 3, 2019 · Apr 1, 2019 · Apr 2, 2019 · Apr 2, 2019 · Apr 2, 2019
diff --git a/bigquery/google/cloud/bigquery/magics.py b/bigquery/google/cloud/bigquery/magics.py
@@ -32,6 +32,11 @@
     * ``--project <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
+        <https://console.cloud.google.com/apis/library/bigquerystorage.googleapis.com>`_.
     * ``--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
+        <https://console.cloud.google.com/apis/library/bigquerystorage.googleapis.com>`_.
+        """
+        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
@@ -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()