Skip to content
This repository was archived by the owner on Mar 6, 2026. It is now read-only.

Commit e8da978

Browse files
LinchinLinchin
authored
chore: refactor BigQueryStorageVersions (#1699)
* chore: refactor BigQueryStorageVersions * address comments in #1680 * add unit test
1 parent 76d0e5a commit e8da978

File tree

13 files changed

+286
-196
lines changed

13 files changed

+286
-196
lines changed

‎google/cloud/bigquery/_helpers.py‎

Lines changed: 0 additions & 65 deletions
Original file line numberDiff line numberDiff line change
@@ -30,10 +30,6 @@
3030
from google.cloud._helpers import _RFC3339_NO_FRACTION
3131
from google.cloud._helpers import _to_bytes
3232

33-
import packaging.version
34-
35-
from google.cloud.bigquery import exceptions
36-
3733
_RFC3339_MICROS_NO_ZULU = "%Y-%m-%dT%H:%M:%S.%f"
3834
_TIMEONLY_WO_MICROS = "%H:%M:%S"
3935
_TIMEONLY_W_MICROS = "%H:%M:%S.%f"
@@ -52,10 +48,6 @@
5248
r"(?P<time_sign>-?)(?P<hours>\d+):(?P<minutes>\d+):(?P<seconds>\d+)\.?(?P<fraction>\d*)?$"
5349
)
5450

55-
_MIN_BQ_STORAGE_VERSION = packaging.version.Version("2.0.0")
56-
57-
_BQ_STORAGE_OPTIONAL_READ_SESSION_VERSION = packaging.version.Version("2.6.0")
58-
5951
BIGQUERY_EMULATOR_HOST = "BIGQUERY_EMULATOR_HOST"
6052
"""Environment variable defining host for emulator."""
6153

@@ -67,63 +59,6 @@ def _get_bigquery_host():
6759
return os.environ.get(BIGQUERY_EMULATOR_HOST, _DEFAULT_HOST)
6860

6961

70-
class BQStorageVersions:
71-
"""Version comparisons for google-cloud-bigqueyr-storage package."""
72-
73-
def __init__(self):
74-
self._installed_version = None
75-
76-
@property
77-
def installed_version(self) -> packaging.version.Version:
78-
"""Return the parsed version of google-cloud-bigquery-storage."""
79-
if self._installed_version is None:
80-
from google.cloud import bigquery_storage
81-
82-
self._installed_version = packaging.version.parse(
83-
# Use 0.0.0, since it is earlier than any released version.
84-
# Legacy versions also have the same property, but
85-
# creating a LegacyVersion has been deprecated.
86-
# https://github.com/pypa/packaging/issues/321
87-
getattr(bigquery_storage, "__version__", "0.0.0")
88-
)
89-
90-
return self._installed_version # type: ignore
91-
92-
@property
93-
def is_read_session_optional(self) -> bool:
94-
"""True if read_session is optional to rows().
95-
96-
See: https://github.com/googleapis/python-bigquery-storage/pull/228
97-
"""
98-
return self.installed_version >= _BQ_STORAGE_OPTIONAL_READ_SESSION_VERSION
99-
100-
def verify_version(self):
101-
"""Verify that a recent enough version of BigQuery Storage extra is
102-
installed.
103-
104-
The function assumes that google-cloud-bigquery-storage extra is
105-
installed, and should thus be used in places where this assumption
106-
holds.
107-
108-
Because `pip` can install an outdated version of this extra despite the
109-
constraints in `setup.py`, the calling code can use this helper to
110-
verify the version compatibility at runtime.
111-
112-
Raises:
113-
exceptions.LegacyBigQueryStorageError:
114-
If the google-cloud-bigquery-storage package is outdated.
115-
"""
116-
if self.installed_version < _MIN_BQ_STORAGE_VERSION:
117-
msg = (
118-
"Dependency google-cloud-bigquery-storage is outdated, please upgrade "
119-
f"it to version >= {_MIN_BQ_STORAGE_VERSION} (version found: {self.installed_version})."
120-
)
121-
raise exceptions.LegacyBigQueryStorageError(msg)
122-
123-
124-
BQ_STORAGE_VERSIONS = BQStorageVersions()
125-
126-
12762
def _not_null(value, field):
12863
"""Check whether 'value' should be coerced to 'field' type."""
12964
return value is not None or (field is not None and field.mode != "NULLABLE")

‎google/cloud/bigquery/_pandas_helpers.py‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,6 @@
2323
import warnings
2424
from typing import Any, Union
2525

26-
from google.cloud.bigquery import _helpers
2726
from google.cloud.bigquery import _pyarrow_helpers
2827
from google.cloud.bigquery import _versions_helpers
2928
from google.cloud.bigquery import schema
@@ -745,7 +744,7 @@ def _download_table_bqstorage_stream(
745744

746745
# Avoid deprecation warnings for passing in unnecessary read session.
747746
# https://github.com/googleapis/python-bigquery-storage/issues/229
748-
if _helpers.BQ_STORAGE_VERSIONS.is_read_session_optional:
747+
if _versions_helpers.BQ_STORAGE_VERSIONS.is_read_session_optional:
749748
rowstream = reader.rows()
750749
else:
751750
rowstream = reader.rows(session)

‎google/cloud/bigquery/_versions_helpers.py‎

Lines changed: 80 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,8 @@
2222

2323

2424
_MIN_PYARROW_VERSION = packaging.version.Version("3.0.0")
25+
_MIN_BQ_STORAGE_VERSION = packaging.version.Version("2.0.0")
26+
_BQ_STORAGE_OPTIONAL_READ_SESSION_VERSION = packaging.version.Version("2.6.0")
2527

2628

2729
class PyarrowVersions:
@@ -51,7 +53,7 @@ def use_compliant_nested_type(self) -> bool:
5153
return self.installed_version.major >= 4
5254

5355
def try_import(self, raise_if_error: bool = False) -> Any:
54-
"""Verify that a recent enough version of pyarrow extra is installed.
56+
"""Verifies that a recent enough version of pyarrow extra is installed.
5557
5658
The function assumes that pyarrow extra is installed, and should thus
5759
be used in places where this assumption holds.
@@ -92,3 +94,80 @@ def try_import(self, raise_if_error: bool = False) -> Any:
9294

9395

9496
PYARROW_VERSIONS = PyarrowVersions()
97+
98+
99+
class BQStorageVersions:
100+
"""Version comparisons for google-cloud-bigqueyr-storage package."""
101+
102+
def __init__(self):
103+
self._installed_version = None
104+
105+
@property
106+
def installed_version(self) -> packaging.version.Version:
107+
"""Return the parsed version of google-cloud-bigquery-storage."""
108+
if self._installed_version is None:
109+
from google.cloud import bigquery_storage
110+
111+
self._installed_version = packaging.version.parse(
112+
# Use 0.0.0, since it is earlier than any released version.
113+
# Legacy versions also have the same property, but
114+
# creating a LegacyVersion has been deprecated.
115+
# https://github.com/pypa/packaging/issues/321
116+
getattr(bigquery_storage, "__version__", "0.0.0")
117+
)
118+
119+
return self._installed_version # type: ignore
120+
121+
@property
122+
def is_read_session_optional(self) -> bool:
123+
"""True if read_session is optional to rows().
124+
125+
See: https://github.com/googleapis/python-bigquery-storage/pull/228
126+
"""
127+
return self.installed_version >= _BQ_STORAGE_OPTIONAL_READ_SESSION_VERSION
128+
129+
def try_import(self, raise_if_error: bool = False) -> Any:
130+
"""Tries to import the bigquery_storage module, and returns results
131+
accordingly. It also verifies the module version is recent enough.
132+
133+
If the import succeeds, returns the ``bigquery_storage`` module.
134+
135+
If the import fails,
136+
returns ``None`` when ``raise_if_error == False``,
137+
raises Error when ``raise_if_error == True``.
138+
139+
Returns:
140+
The ``bigquery_storage`` module or ``None``.
141+
142+
Raises:
143+
exceptions.BigQueryStorageNotFoundError:
144+
If google-cloud-bigquery-storage is not installed
145+
exceptions.LegacyBigQueryStorageError:
146+
If google-cloud-bigquery-storage package is outdated
147+
"""
148+
try:
149+
from google.cloud import bigquery_storage # type: ignore
150+
except ImportError:
151+
if raise_if_error:
152+
msg = (
153+
"Package google-cloud-bigquery-storage not found. "
154+
"Install google-cloud-bigquery-storage version >= "
155+
f"{_MIN_BQ_STORAGE_VERSION}."
156+
)
157+
raise exceptions.BigQueryStorageNotFoundError(msg)
158+
return None
159+
160+
if self.installed_version < _MIN_BQ_STORAGE_VERSION:
161+
if raise_if_error:
162+
msg = (
163+
"Dependency google-cloud-bigquery-storage is outdated, "
164+
f"please upgrade it to version >= {_MIN_BQ_STORAGE_VERSION} "
165+
f"(version found: {self.installed_version})."
166+
)
167+
raise exceptions.LegacyBigQueryStorageError(msg)
168+
return None
169+
170+
return bigquery_storage
171+
172+
173+
BQ_STORAGE_VERSIONS = BQStorageVersions()

‎google/cloud/bigquery/client.py‎

Lines changed: 22 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -65,26 +65,25 @@
6565
DEFAULT_BQSTORAGE_CLIENT_INFO = None # type: ignore
6666

6767

68+
from google.cloud.bigquery._http import Connection
6869
from google.cloud.bigquery import _job_helpers
69-
from google.cloud.bigquery._job_helpers import make_job_id as _make_job_id
70+
from google.cloud.bigquery import _pandas_helpers
71+
from google.cloud.bigquery import _versions_helpers
72+
from google.cloud.bigquery import enums
73+
from google.cloud.bigquery import exceptions as bq_exceptions
74+
from google.cloud.bigquery import job
7075
from google.cloud.bigquery._helpers import _get_sub_prop
7176
from google.cloud.bigquery._helpers import _record_field_to_json
7277
from google.cloud.bigquery._helpers import _str_or_none
7378
from google.cloud.bigquery._helpers import _verify_job_config_type
7479
from google.cloud.bigquery._helpers import _get_bigquery_host
75-
from google.cloud.bigquery._helpers import BQ_STORAGE_VERSIONS
7680
from google.cloud.bigquery._helpers import _DEFAULT_HOST
77-
from google.cloud.bigquery._http import Connection
78-
from google.cloud.bigquery import _pandas_helpers
79-
from google.cloud.bigquery import _versions_helpers
81+
from google.cloud.bigquery._job_helpers import make_job_id as _make_job_id
8082
from google.cloud.bigquery.dataset import Dataset
8183
from google.cloud.bigquery.dataset import DatasetListItem
8284
from google.cloud.bigquery.dataset import DatasetReference
83-
from google.cloud.bigquery import enums
8485
from google.cloud.bigquery.enums import AutoRowIDs
85-
from google.cloud.bigquery import exceptions as bq_exceptions
86-
from google.cloud.bigquery.opentelemetry_tracing import create_span
87-
from google.cloud.bigquery import job
86+
from google.cloud.bigquery.format_options import ParquetOptions
8887
from google.cloud.bigquery.job import (
8988
CopyJob,
9089
CopyJobConfig,
@@ -98,6 +97,7 @@
9897
from google.cloud.bigquery.model import Model
9998
from google.cloud.bigquery.model import ModelReference
10099
from google.cloud.bigquery.model import _model_arg_to_model_ref
100+
from google.cloud.bigquery.opentelemetry_tracing import create_span
101101
from google.cloud.bigquery.query import _QueryResults
102102
from google.cloud.bigquery.retry import (
103103
DEFAULT_JOB_RETRY,
@@ -113,7 +113,6 @@
113113
from google.cloud.bigquery.table import TableListItem
114114
from google.cloud.bigquery.table import TableReference
115115
from google.cloud.bigquery.table import RowIterator
116-
from google.cloud.bigquery.format_options import ParquetOptions
117116

118117
pyarrow = _versions_helpers.PYARROW_VERSIONS.try_import()
119118

@@ -545,29 +544,32 @@ def _ensure_bqstorage_client(
545544
An existing BigQuery Storage client instance. If ``None``, a new
546545
instance is created and returned.
547546
client_options:
548-
Custom options used with a new BigQuery Storage client instance if one
549-
is created.
547+
Custom options used with a new BigQuery Storage client instance
548+
if one is created.
550549
client_info:
551-
The client info used with a new BigQuery Storage client instance if one
552-
is created.
550+
The client info used with a new BigQuery Storage client
551+
instance if one is created.
553552
554553
Returns:
555554
A BigQuery Storage API client.
556555
"""
556+
557557
try:
558-
from google.cloud import bigquery_storage # type: ignore
559-
except ImportError:
558+
bigquery_storage = _versions_helpers.BQ_STORAGE_VERSIONS.try_import(
559+
raise_if_error=True
560+
)
561+
except bq_exceptions.BigQueryStorageNotFoundError:
560562
warnings.warn(
561563
"Cannot create BigQuery Storage client, the dependency "
562564
"google-cloud-bigquery-storage is not installed."
563565
)
564566
return None
565-
566-
try:
567-
BQ_STORAGE_VERSIONS.verify_version()
568567
except bq_exceptions.LegacyBigQueryStorageError as exc:
569-
warnings.warn(str(exc))
568+
warnings.warn(
569+
"Dependency google-cloud-bigquery-storage is outdated: " + str(exc)
570+
)
570571
return None
572+
571573
if bqstorage_client is None:
572574
bqstorage_client = bigquery_storage.BigQueryReadClient(
573575
credentials=self._credentials,

‎google/cloud/bigquery/exceptions.py‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,3 +23,9 @@ class LegacyBigQueryStorageError(BigQueryError):
2323

2424
class LegacyPyarrowError(BigQueryError):
2525
"""Raised when too old a version of pyarrow package is detected at runtime."""
26+
27+
28+
class BigQueryStorageNotFoundError(BigQueryError):
29+
"""Raised when BigQuery Storage extra is not installed when trying to
30+
import it.
31+
"""

‎google/cloud/bigquery/magics/magics.py‎

Lines changed: 25 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -104,6 +104,8 @@
104104
import google.auth # type: ignore
105105
from google.cloud import bigquery
106106
import google.cloud.bigquery.dataset
107+
from google.cloud.bigquery import _versions_helpers
108+
from google.cloud.bigquery import exceptions
107109
from google.cloud.bigquery.dbapi import _helpers
108110
from google.cloud.bigquery.magics import line_arg_parser as lap
109111

@@ -744,19 +746,40 @@ def _split_args_line(line):
744746

745747

746748
def _make_bqstorage_client(client, use_bqstorage_api, client_options):
749+
"""Creates a BigQuery Storage client.
750+
751+
Args:
752+
client (:class:`~google.cloud.bigquery.client.Client`): BigQuery client.
753+
use_bqstorage_api (bool): whether BigQuery Storage API is used or not.
754+
client_options (:class:`google.api_core.client_options.ClientOptions`):
755+
Custom options used with a new BigQuery Storage client instance
756+
if one is created.
757+
758+
Raises:
759+
ImportError: if google-cloud-bigquery-storage is not installed, or
760+
grpcio package is not installed.
761+
762+
763+
Returns:
764+
None: if ``use_bqstorage_api == False``, or google-cloud-bigquery-storage
765+
is outdated.
766+
BigQuery Storage Client:
767+
"""
747768
if not use_bqstorage_api:
748769
return None
749770

750771
try:
751-
from google.cloud import bigquery_storage # type: ignore # noqa: F401
752-
except ImportError as err:
772+
_versions_helpers.BQ_STORAGE_VERSIONS.try_import(raise_if_error=True)
773+
except exceptions.BigQueryStorageNotFoundError as err:
753774
customized_error = ImportError(
754775
"The default BigQuery Storage API client cannot be used, install "
755776
"the missing google-cloud-bigquery-storage and pyarrow packages "
756777
"to use it. Alternatively, use the classic REST API by specifying "
757778
"the --use_rest_api magic option."
758779
)
759780
raise customized_error from err
781+
except exceptions.LegacyBigQueryStorageError:
782+
pass
760783

761784
try:
762785
from google.api_core.gapic_v1 import client_info as gapic_client_info

0 commit comments

Comments
 (0)