From 55d7dc6d204c0634b9cffbe4c181099f1910b609 Mon Sep 17 00:00:00 2001 From: Tim Swast Date: Wed, 17 Apr 2019 15:01:08 -0700 Subject: [PATCH] Raise ValueError when BQ Storage is required but missing Move the BigQuery Storage API optional dependency imports to the top of the module, following the existing pattern for other optional dependencies. --- bigquery/google/cloud/bigquery/table.py | 30 +++++++++++++++++--- bigquery/tests/unit/test_table.py | 37 +++++++++++++++++++++++++ 2 files changed, 63 insertions(+), 4 deletions(-) diff --git a/bigquery/google/cloud/bigquery/table.py b/bigquery/google/cloud/bigquery/table.py index d50b43a18058..92af19c43ce5 100644 --- 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_dataframe() 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_dataframe_tabledata_list(self, dtypes, progress_bar=None): def _to_dataframe_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_dataframe(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 index e8ebbdfc6c90..7ac982394c9d 100644 --- a/bigquery/tests/unit/test_table.py +++ b/bigquery/tests/unit/test_table.py @@ -1860,6 +1860,28 @@ def test_to_dataframe_w_bqstorage_raises_auth_error(self): with pytest.raises(google.api_core.exceptions.Forbidden): row_iterator.to_dataframe(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_dataframe_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_dataframe(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)