chore: clean up BigQueryStorage import checks

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."""

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."

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.
+    """

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 "

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:

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):