8000 feat: detect obsolete BQ Storage extra at runtime (#666) · googleapis/python-bigquery@bd7dbda · GitHub
[go: up one dir, main page]
Skip to content

Commit bd7dbda

Browse files
plamutplamut
authored
feat: detect obsolete BQ Storage extra at runtime (#666)
* feat: detect obsolete BQ Storage extra at runtime * Cover the changes with unit tests * Skip BQ Storage version tests if extra missing * Rename and improve _create_bqstorage_client() The method is renamed to _ensure_bqstorage_client() and now performs a check if BQ Storage dependency is recent enough. * Remove BQ Storage check from dbapi.Cursor The check is now performed in dbapi.Connection, which is sufficient. * Remove BQ Storage check in _pandas_helpers The methods in higher layers already do the same check before a BQ Storage client instance is passed to _pandas_helpers._download_table_bqstorage() helper. * Simplify BQ Storage client factory in magics Lean more heavily on client._ensure_bqstorage_client() to de-duplicate logic. * Cover missing code lines with tests
1 parent 82f6c32 commit bd7dbda

File tree

13 files changed

+357
-36
lines changed

13 files changed

+357
-36
lines changed

‎google/cloud/bigquery/__init__.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,7 @@
3939
from google.cloud.bigquery import enums
4040
from google.cloud.bigquery.enums import SqlTypeNames
4141
from google.cloud.bigquery.enums import StandardSqlDataTypes
42+
from google.cloud.bigquery.exceptions import LegacyBigQueryStorageError
4243
from google.cloud.bigquery.external_config import ExternalConfig
4344
from google.cloud.bigquery.external_config import BigtableOptions
4445
from google.cloud.bigquery.external_config import BigtableColumnFamily
@@ -152,6 +153,8 @@
152153
"WriteDisposition",
153154
# EncryptionConfiguration
154155
"EncryptionConfiguration",
156+
# Custom exceptions
157+
"LegacyBigQueryStorageError",
155158
]
156159

157160

‎google/cloud/bigquery/_helpers.py

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,10 @@
2525
from google.cloud._helpers import _RFC3339_MICROS
2626
from google.cloud._helpers import _RFC3339_NO_FRACTION
2727
from google.cloud._helpers import _to_bytes
28+
import pkg_resources
29+
30+
from google.cloud.bigquery.exceptions import LegacyBigQueryStorageError
31+
2832

2933
_RFC3339_MICROS_NO_ZULU = "%Y-%m-%dT%H:%M:%S.%f"
3034
_TIMEONLY_WO_MICROS = "%H:%M:%S"
@@ -36,6 +40,32 @@
3640
re.VERBOSE,
3741
)
3842

43+
_MIN_BQ_STORAGE_VERSION = pkg_resources.parse_version("2.0.0")
44+
45+
46+
def _verify_bq_storage_version():
47+
"""Verify that a recent enough version of BigQuery Storage extra is installed.
48+
49+
The function assumes that google-cloud-bigquery-storage extra is installed, and
50+
should thus be used in places where this assumption holds.
51+
52+
Because `pip` can install an outdated version of this extra despite the constraints
53+
in setup.py, the the calling code can use this helper to verify the version
54+
compatibility at runtime.
55+
"""
56+
from google.cloud import bigquery_storage
57+
58+
installed_version = pkg_resources.parse_version(
59+
getattr(bigquery_storage, "__version__", "legacy")
60+
)
61+
62+
if installed_version < _MIN_BQ_STORAGE_VERSION:
63+
msg = (
64+
"Dependency google-cloud-bigquery-storage is outdated, please upgrade "
65+
f"it to version >= 2.0.0 (version found: {installed_version})."
66+
)
67+
raise LegacyBigQueryStorageError(msg)
68+
3969

4070
def _not_null(value, field):
4171
"""Check whether 'value' should be coerced to 'field' type."""

‎google/cloud/bigquery/client.py

Lines changed: 51 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -50,16 +50,25 @@
5050
from google.cloud import exceptions # pytype: disable=import-error
5151
from google.cloud.client import ClientWithProject # pytype: disable=import-error
5252

53+
try:
54+
from google.cloud.bigquery_storage_v1.services.big_query_read.client import (
55+
DEFAULT_CLIENT_INFO as DEFAULT_BQSTORAGE_CLIENT_INFO,
56+
)
57+
except ImportError:
58+
DEFAULT_BQSTORAGE_CLIENT_INFO = None
59+
5360
from google.cloud.bigquery._helpers import _del_sub_prop
5461
from google.cloud.bigquery._helpers import _get_sub_prop
5562
from google.cloud.bigquery._helpers import _record_field_to_json
5663
from google.cloud.bigquery._helpers import _str_or_none
64+
from google.cloud.bigquery._helpers import _verify_bq_storage_version
5765
from google.cloud.bigquery._helpers import _verify_job_config_type
5866
from google.cloud.bigquery._http import Connection
5967
from google.cloud.bigquery import _pandas_helpers
6068
from google.cloud.bigquery.dataset import Dataset
6169
from google.cloud.bigquery.dataset import DatasetListItem
6270
from google.cloud.bigquery.dataset import DatasetReference
71+
from google.cloud.bigquery.exceptions import LegacyBigQueryStorageError
6372
from google.cloud.bigquery.opentelemetry_tracing import create_span
6473
from google.cloud.bigquery import job
6574
from google.cloud.bigquery.job import (
@@ -445,15 +454,38 @@ def dataset(self, dataset_id: str, project: str = None) -> DatasetReference:
445454
)
446455
return DatasetReference(project, dataset_id)
447456

448-
def _create_bqstorage_client(self):
457+
def _ensure_bqstorage_client(
458+
self,
459+
bqstorage_client: Optional[
460+
"google.cloud.bigquery_storage.BigQueryReadClient"
461+
] = None,
462+
client_options: Optional[google.api_core.client_options.ClientOptions] = None,
463+
client_info: Optional[
464+
"google.api_core.gapic_v1.client_info.ClientInfo"
465+
] = DEFAULT_BQSTORAGE_CLIENT_INFO,
466+
) -> Optional["google.cloud.bigquery_storage.BigQueryReadClient"]:
449467
"""Create a BigQuery Storage API client using this client's credentials.
450468
451-
If a client cannot be created due to missing dependencies, raise a
452-
warning and return ``None``.
469+
If a client cannot be created due to a missing or outdated dependency
470+
`google-cloud-bigquery-storage`, raise a warning and return ``None``.
471+
472+
If the `bqstorage_client` argument is not ``None``, still perform the version
473+
check and return the argument back to the caller if the check passes. If it
474+
fails, raise a warning and return ``None``.
475+
476+
Args:
477+
bqstorage_client:
478+
An existing BigQuery Storage client instance to check for version
479+
compatibility. If ``None``, a new instance is created and returned.
480+
client_options:
481+
Custom options used with a new BigQuery Storage client instance if one
482+
is created.
483+
client_info:
484+
The client info used with a new BigQuery Storage client instance if one
485+
is created.
453486
454487
Returns:
455-
Optional[google.cloud.bigquery_storage.BigQueryReadClient]:
456-
A BigQuery Storage API client.
488+
A BigQuery Storage API client.
457489
"""
458490
try:
459491
from google.cloud import bigquery_storage
@@ -464,7 +496,20 @@ def _create_bqstorage_client(self):
464496
)
465497
return None
466498

467-
return bigquery_storage.BigQueryReadClient(credentials=self._credentials)
499+
try:
500+
_verify_bq_storage_version()
501+
except LegacyBigQueryStorageError as exc:
502+
warnings.warn(str(exc))
503+
return None
504+
505+
if bqstorage_client is None:
506+
bqstorage_client = bigquery_storage.BigQueryReadClient(
507+
credentials=self._credentials,
508+
client_options=client_options,
509+
client_info=client_info,
510+
)
511+
512+
return bqstorage_client
468513

469514
def _dataset_from_arg(self, dataset):
470515
if isinstance(dataset, str):

‎google/cloud/bigquery/dbapi/connection.py

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -47,12 +47,14 @@ def __init__(self, client=None, bqstorage_client=None):
4747
else:
4848
self._owns_client = False
4949

50+
# A warning is already raised by the BQ Storage client factory factory if
51+
# instantiation fails, or if the given BQ Storage client instance is outdated.
5052
if bqstorage_client is None:
51-
# A warning is already raised by the factory if instantiation fails.
52-
bqstorage_client = client._create_bqstorage_client()
53+
bqstorage_client = client._ensure_bqstorage_client()
5354
self._owns_bqstorage_client = bqstorage_client is not None
5455
else:
5556
self._owns_bqstorage_client = False
57+
bqstorage_client = client._ensure_bqstorage_client(bqstorage_client)
5658

5759
self._client = client
5860
self._bqstorage_client = bqstorage_client

‎google/cloud/bigquery/exceptions.py

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,21 @@
1+
# Copyright 2021 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
16+
class BigQueryError(Exception):
17+
"""Base class for all custom exceptions defined by the BigQuery client."""
18+
19+
20+
class LegacyBigQueryStorageError(BigQueryError):
21+
"""Raised when too old a version of BigQuery Storage extra is detected at runtime."""

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

Lines changed: 5 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -644,7 +644,7 @@ def _cell_magic(line, query):
644644
bqstorage_client_options.api_endpoint = args.bqstorage_api_endpoint
645645

646646
bqstorage_client = _make_bqstorage_client(
647-
use_bqstorage_api, context.credentials, bqstorage_client_options,
647+
client, use_bqstorage_api, bqstorage_client_options,
648648
)
649649

650650
close_transports = functools.partial(_close_transports, client, bqstorage_client)
@@ -762,12 +762,12 @@ def _split_args_line(line):
762762
return params_option_value, rest_of_args
763763

764764

765-
def _make_bqstorage_client(use_bqstorage_api, credentials, client_options):
765+
def _make_bqstorage_client(client, use_bqstorage_api, client_options):
766766
if not use_bqstorage_api:
767767
return None
768768

769769
try:
770-
from google.cloud import bigquery_storage
770+
from google.cloud import bigquery_storage # noqa: F401
771771
except ImportError as err:
772772
customized_error = ImportError(
773773
"The default BigQuery Storage API client cannot be used, install "
@@ -785,10 +785,9 @@ def _make_bqstorage_client(use_bqstorage_api, credentials, client_options):
785785
)
786786
raise customized_error from err
787787

788-
return bigquery_storage.BigQueryReadClient(
789-
credentials=credentials,
790-
client_info=gapic_client_info.ClientInfo(user_agent=IPYTHON_USER_AGENT),
788+
return client._ensure_bqstorage_client(
791789
client_options=client_options,
790+
client_info=gapic_client_info.ClientInfo(user_agent=IPYTHON_USER_AGENT),
792791
)
793792

794793

‎google/cloud/bigquery/table.py

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,7 @@
4141
import google.cloud._helpers
4242
from google.cloud.bigquery import _helpers
4343
from google.cloud.bigquery import _pandas_helpers
44+
from google.cloud.bigquery.exceptions import LegacyBigQueryStorageError
4445
from google.cloud.bigquery.schema import _build_schema_resource
4546
from google.cloud.bigquery.schema import _parse_schema_resource
4647
from google.cloud.bigquery.schema import _to_schema_fields
@@ -1519,6 +1520,17 @@ def _validate_bqstorage(self, bqstorage_client, create_bqstorage_client):
15191520
)
15201521
return False
15211522

1523+
try:
1524+
from google.cloud import bigquery_storage # noqa: F401
1525+
except ImportError:
1526+
return False
1527+
1528+
try:
1529+
_helpers._verify_bq_storage_version()
1530+
except LegacyBigQueryStorageError as exc:
1531+
warnings.warn(str(exc))
1532+
return False
1533+
15221534
return True
15231535

15241536
def _get_next_page_response(self):
@@ -1655,7 +1667,7 @@ def to_arrow(
16551667

16561668
owns_bqstorage_client = False
16571669
if not bqstorage_client and create_bqstorage_client:
1658-
bqstorage_client = self.client._create_bqstorage_client()
1670+
bqstorage_client = self.client._ensure_bqstorage_client()
16591671
owns_bqstorage_client = bqstorage_client is not None
16601672

16611673
try:

‎tests/unit/test__helpers.py

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,44 @@
1919

2020
import mock
2121

22+
try:
23+
from google.cloud import bigquery_storage
24+
except ImportError: # pragma: NO COVER
25+
bigquery_storage = None
26+
27+
28+
@unittest.skipIf(bigquery_storage is None, "Requires `google-cloud-bigquery-storage`")
29+
class Test_verify_bq_storage_version(unittest.TestCase):
30+
def _call_fut(self):
31+
from google.cloud.bigquery._helpers import _verify_bq_storage_version
32+
33+
return _verify_bq_storage_version()
34+
35+
def test_raises_no_error_w_recent_bqstorage(self):
36+
from google.cloud.bigquery.exceptions import LegacyBigQueryStorageError
37+
38+
with mock.patch("google.cloud.bigquery_storage.__version__", new="2.0.0"):
39+
try:
40+
self._call_fut()
41+
except LegacyBigQueryStorageError: # pragma: NO COVER
42+
self.fail("Legacy error raised with a non-legacy dependency version.")
43+
44+
def test_raises_error_w_legacy_bqstorage(self):
45+
from google.cloud.bigquery.exceptions import LegacyBigQueryStorageError
46+
47+
with mock.patch("google.cloud.bigquery_storage.__version__", new="1.9.9"):
48+
with self.assertRaises(LegacyBigQueryStorageError):
49+
self._call_fut()
50+
51+
def test_raises_error_w_unknown_bqstorage_version(self):
52+
from google.cloud.bigquery.exceptions import LegacyBigQueryStorageError
53+
54+
with mock.patch("google.cloud.bigquery_storage", autospec=True) as fake_module:
55+
del fake_module.__version__
56+
error_pattern = r"version found: legacy"
57+
with self.assertRaisesRegex(LegacyBigQueryStorageError, error_pattern):
58+
self._call_fut()
59+
2260

2361
class Test_not_null(unittest.TestCase):
2462
def _call_fut(self, value, field):

0 commit comments

Comments
 (0)
0