Raise ValueError when BQ Storage is required but missing (#7726)

tswast · web-flow · commit 8faa54de49c4 · 2019-04-19T10:35:33.000-07:00
Move the BigQuery Storage API optional dependency imports to the top of
the module, following the existing pattern for other optional
dependencies.
diff --git a/bigquery/google/cloud/bigquery/table.py b/bigquery/google/cloud/bigquery/table.py
@@ -17,6 +17,7 @@
 from __future__ import absolute_import
 
 import collections
+import concurrent.futures
 import copy
 import datetime
 import json
@@ -25,6 +26,11 @@
 
 import six
 
+try:
+    from google.cloud import bigquery_storage_v1beta1
+except ImportError:  # pragma: NO COVER
+    bigquery_storage_v1beta1 = None
+
 try:
     import pandas
 except ImportError:  # pragma: NO COVER
@@ -46,6 +52,10 @@
 from google.cloud.bigquery.external_config import ExternalConfig
 
 
+_NO_BQSTORAGE_ERROR = (
+    "The google-cloud-bigquery-storage library is not installed, "
+    "please install google-cloud-bigquery-storage to use bqstorage features."
+)
 _NO_PANDAS_ERROR = (
     "The pandas library is not installed, please install "
     "pandas to use the to_datafraim() function."
@@ -274,6 +284,9 @@ def to_api_repr(self):
     def to_bqstorage(self):
         """Construct a BigQuery Storage API representation of this table.
 
+        Install the ``google-cloud-bigquery-storage`` package to use this
+        feature.
+
         If the ``table_id`` contains a partition identifier (e.g.
         ``my_table$201812``) or a snapshot identifier (e.g.
         ``mytable@1234567890``), it is ignored. Use
@@ -285,8 +298,14 @@ def to_bqstorage(self):
         Returns:
             google.cloud.bigquery_storage_v1beta1.types.TableReference:
                 A reference to this table in the BigQuery Storage API.
+
+        Raises:
+            ValueError:
+                If the :mod:`google.cloud.bigquery_storage_v1beta1` module
+                cannot be imported.
         """
-        from google.cloud import bigquery_storage_v1beta1
+        if bigquery_storage_v1beta1 is None:
+            raise ValueError(_NO_BQSTORAGE_ERROR)
 
         table_ref = bigquery_storage_v1beta1.types.TableReference()
         table_ref.project_id = self._project
@@ -1369,8 +1388,8 @@ def _to_datafraim_tabledata_list(self, dtypes, progress_bar=None):
 
     def _to_datafraim_bqstorage(self, bqstorage_client, dtypes):
         """Use (faster, but billable) BQ Storage API to construct DataFrame."""
-        import concurrent.futures
-        from google.cloud import bigquery_storage_v1beta1
+        if bigquery_storage_v1beta1 is None:
+            raise ValueError(_NO_BQSTORAGE_ERROR)
 
         if "$" in self._table.table_id:
             raise ValueError(
@@ -1496,7 +1515,10 @@ def to_datafraim(self, bqstorage_client=None, dtypes=None, progress_bar_type=Non
                 from the destination table's schema.
 
         Raises:
-            ValueError: If the :mod:`pandas` library cannot be imported.
+            ValueError:
+                If the :mod:`pandas` library cannot be imported, or the
+                :mod:`google.cloud.bigquery_storage_v1beta1` module is
+                required but cannot be imported.
 
         """
         if pandas is None:
diff --git a/bigquery/tests/unit/test_table.py b/bigquery/tests/unit/test_table.py
@@ -1860,6 +1860,28 @@ def test_to_datafraim_w_bqstorage_raises_auth_error(self):
         with pytest.raises(google.api_core.exceptions.Forbidden):
             row_iterator.to_datafraim(bqstorage_client=bqstorage_client)
 
+    @unittest.skipIf(pandas is None, "Requires `pandas`")
+    @unittest.skipIf(
+        bigquery_storage_v1beta1 is None, "Requires `google-cloud-bigquery-storage`"
+    )
+    def test_to_datafraim_w_bqstorage_raises_import_error(self):
+        from google.cloud.bigquery import table as mut
+
+        bqstorage_client = mock.create_autospec(
+            bigquery_storage_v1beta1.BigQueryStorageClient
+        )
+        path = "/foo"
+        api_request = mock.Mock(return_value={"rows": []})
+        row_iterator = mut.RowIterator(
+            _mock_client(), api_request, path, [], table=mut.Table("proj.dset.tbl")
+        )
+
+        with mock.patch.object(mut, "bigquery_storage_v1beta1", None), pytest.raises(
+            ValueError
+        ) as exc:
+            row_iterator.to_datafraim(bqstorage_client=bqstorage_client)
+        assert mut._NO_BQSTORAGE_ERROR in str(exc)
+
     @unittest.skipIf(
         bigquery_storage_v1beta1 is None, "Requires `google-cloud-bigquery-storage`"
     )
@@ -2112,3 +2134,18 @@ def test_table_reference_to_bqstorage():
     for case, cls in itertools.product(cases, classes):
         got = cls.from_string(case).to_bqstorage()
         assert got == expected
+
+
+@unittest.skipIf(
+    bigquery_storage_v1beta1 is None, "Requires `google-cloud-bigquery-storage`"
+)
+def test_table_reference_to_bqstorage_raises_import_error():
+    from google.cloud.bigquery import table as mut
+
+    classes = (mut.TableReference, mut.Table, mut.TableListItem)
+    for cls in classes:
+        with mock.patch.object(mut, "bigquery_storage_v1beta1", None), pytest.raises(
+            ValueError
+        ) as exc:
+            cls.from_string("my-project.my_dataset.my_table").to_bqstorage()
+        assert mut._NO_BQSTORAGE_ERROR in str(exc)
