diff --git a/google/cloud/bigquery/_helpers.py b/google/cloud/bigquery/_helpers.py index 014a721a8..e98c71249 100644 --- a/google/cloud/bigquery/_helpers.py +++ b/google/cloud/bigquery/_helpers.py @@ -35,6 +35,7 @@ from google.cloud.bigquery.exceptions import ( LegacyBigQueryStorageError, LegacyPyarrowError, + BigQueryStorageNotFoundError, ) _RFC3339_MICROS_NO_ZULU = "%Y-%m-%dT%H:%M:%S.%f" @@ -120,11 +121,34 @@ def verify_version(self): """ if self.installed_version < _MIN_BQ_STORAGE_VERSION: msg = ( - "Dependency google-cloud-bigquery-storage is outdated, please upgrade " - f"it to version >= {_MIN_BQ_STORAGE_VERSION} (version found: {self.installed_version})." + "Dependency google-cloud-bigquery-storage is outdated, " + f"please upgrade it to version >= {_MIN_BQ_STORAGE_VERSION} " + f"(version found: {self.installed_version})." ) raise LegacyBigQueryStorageError(msg) + def try_import(self) -> Any: + """Tries to import the bigquery_storage module, and returns an + error if BigQuery Storage extra is not installed. + + Returns: + The ``bigquery_storage`` module. + + Raises: + BigQueryStorageNotFoundError: + If google-cloud-bigquery-storage is not installed. + """ + try: + from google.cloud import bigquery_storage # type: ignore + except ImportError: + msg = ( + "Package google-cloud-bigquery-storage not found. " + "Install google-cloud-bigquery-storage version >= " + f"{_MIN_BQ_STORAGE_VERSION}." + ) + raise BigQueryStorageNotFoundError(msg) + return bigquery_storage + class PyarrowVersions: """Version comparisons for pyarrow package.""" diff --git a/google/cloud/bigquery/client.py b/google/cloud/bigquery/client.py index f7c7864a1..e1f585c42 100644 --- a/google/cloud/bigquery/client.py +++ b/google/cloud/bigquery/client.py @@ -89,7 +89,10 @@ from google.cloud.bigquery.dataset import DatasetReference from google.cloud.bigquery import enums from google.cloud.bigquery.enums import AutoRowIDs -from google.cloud.bigquery.exceptions import LegacyBigQueryStorageError +from google.cloud.bigquery.exceptions import ( + LegacyBigQueryStorageError, + BigQueryStorageNotFoundError, +) from google.cloud.bigquery.opentelemetry_tracing import create_span from google.cloud.bigquery import job from google.cloud.bigquery.job import ( @@ -564,8 +567,8 @@ def _ensure_bqstorage_client( A BigQuery Storage API client. """ try: - from google.cloud import bigquery_storage # type: ignore - except ImportError: + bigquery_storage = BQ_STORAGE_VERSIONS.try_import() + except BigQueryStorageNotFoundError: warnings.warn( "Cannot create BigQuery Storage client, the dependency " "google-cloud-bigquery-storage is not installed." diff --git a/google/cloud/bigquery/exceptions.py b/google/cloud/bigquery/exceptions.py index 2bab97fea..e94a6c832 100644 --- a/google/cloud/bigquery/exceptions.py +++ b/google/cloud/bigquery/exceptions.py @@ -23,3 +23,9 @@ class LegacyBigQueryStorageError(BigQueryError): class LegacyPyarrowError(BigQueryError): """Raised when too old a version of pyarrow package is detected at runtime.""" + + +class BigQueryStorageNotFoundError(BigQueryError): + """Raised when BigQuery Storage extra is not installed when trying to + import it. + """ diff --git a/google/cloud/bigquery/magics/magics.py b/google/cloud/bigquery/magics/magics.py index f92f77541..7bb8e791c 100644 --- a/google/cloud/bigquery/magics/magics.py +++ b/google/cloud/bigquery/magics/magics.py @@ -103,8 +103,10 @@ from google.api_core.exceptions import NotFound import google.auth # type: ignore from google.cloud import bigquery +from google.cloud.bigquery._helpers import BQ_STORAGE_VERSIONS import google.cloud.bigquery.dataset from google.cloud.bigquery.dbapi import _helpers +from google.cloud.bigquery.exceptions import BigQueryStorageNotFoundError from google.cloud.bigquery.magics import line_arg_parser as lap @@ -748,8 +750,8 @@ def _make_bqstorage_client(client, use_bqstorage_api, client_options): return None try: - from google.cloud import bigquery_storage # type: ignore # noqa: F401 - except ImportError as err: + BQ_STORAGE_VERSIONS.try_import() # type: ignore # noqa: F401 + except BigQueryStorageNotFoundError as err: customized_error = ImportError( "The default BigQuery Storage API client cannot be used, install " "the missing google-cloud-bigquery-storage and pyarrow packages " diff --git a/google/cloud/bigquery/table.py b/google/cloud/bigquery/table.py index 462447d51..957c502b0 100644 --- a/google/cloud/bigquery/table.py +++ b/google/cloud/bigquery/table.py @@ -61,7 +61,10 @@ from google.cloud.bigquery import _helpers from google.cloud.bigquery import _pandas_helpers from google.cloud.bigquery.enums import DefaultPandasDTypes -from google.cloud.bigquery.exceptions import LegacyBigQueryStorageError +from google.cloud.bigquery.exceptions import ( + LegacyBigQueryStorageError, + BigQueryStorageNotFoundError, +) from google.cloud.bigquery.schema import _build_schema_resource from google.cloud.bigquery.schema import _parse_schema_resource from google.cloud.bigquery.schema import _to_schema_fields @@ -1610,8 +1613,8 @@ def _validate_bqstorage(self, bqstorage_client, create_bqstorage_client): return False try: - from google.cloud import bigquery_storage # noqa: F401 - except ImportError: + _helpers.BQ_STORAGE_VERSIONS.try_import() + except BigQueryStorageNotFoundError: return False try: diff --git a/tests/unit/test__helpers.py b/tests/unit/test__helpers.py index 4fb86f665..b53b66904 100644 --- a/tests/unit/test__helpers.py +++ b/tests/unit/test__helpers.py @@ -49,6 +49,12 @@ def _call_fut(self): _helpers.BQ_STORAGE_VERSIONS._installed_version = None return _helpers.BQ_STORAGE_VERSIONS.verify_version() + def _call_try_import(self): + from google.cloud.bigquery import _helpers + + _helpers.BQ_STORAGE_VERSIONS._installed_version = None + return _helpers.BQ_STORAGE_VERSIONS.try_import() + def test_raises_no_error_w_recent_bqstorage(self): from google.cloud.bigquery.exceptions import LegacyBigQueryStorageError @@ -99,6 +105,9 @@ def test_is_read_session_optional_false(self): with mock.patch("google.cloud.bigquery_storage.__version__", new="2.5.0"): assert not versions.is_read_session_optional + def test_try_import(self): + assert self._call_try_import() is not None + @unittest.skipIf(pyarrow is None, "Requires `pyarrow`") class TestPyarrowVersions(unittest.TestCase):