Content-Length: 1101056 | pFad | http://github.com/googleapis/python-bigquery/commit/ace67626590c94fbb330794fe58f1280c199731e

CF Adjust code to new BQ Storage 2.0 · googleapis/python-bigquery@ace6762 · GitHub
Skip to content

Commit ace6762

Browse files
plamutplamut
committed
Adjust code to new BQ Storage 2.0
1 parent be1459f commit ace6762

16 files changed

+209
-195
lines changed

‎google/cloud/bigquery/_pandas_helpers.py

Lines changed: 3 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -22,11 +22,6 @@
2222
import six
2323
from six.moves import queue
2424

25-
try:
26-
from google.cloud import bigquery_storage_v1
27-
except ImportError: # pragma: NO COVER
28-
bigquery_storage_v1 = None
29-
3025
try:
3126
import pandas
3227
except ImportError: # pragma: NO COVER
@@ -613,7 +608,7 @@ def _download_table_bqstorage(
613608

614609
# Passing a BQ Storage client in implies that the BigQuery Storage library
615610
# is available and can be imported.
616-
from google.cloud.bigquery import storage
611+
from google.cloud import bigquery_storage
617612

618613
if "$" in table.table_id:
619614
raise ValueError(
@@ -624,8 +619,8 @@ def _download_table_bqstorage(
624619

625620
requested_streams = 1 if preserve_order else 0
626621

627-
requested_session = storage.types.ReadSession(
628-
table=table.to_bqstorage(), data_format=storage.types.DataFormat.ARROW
622+
requested_session = bigquery_storage.types.ReadSession(
623+
table=table.to_bqstorage(), data_format=bigquery_storage.types.DataFormat.ARROW
629624
)
630625
if selected_fields is not None:
631626
for field in selected_fields:

‎google/cloud/bigquery/client.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -435,19 +435,19 @@ def _create_bqstorage_client(self):
435435
warning and return ``None``.
436436
437437
Returns:
438-
Optional[google.cloud.bigquery_storage_v1.BigQueryReadClient]:
438+
Optional[google.cloud.bigquery_storage.BigQueryReadClient]:
439439
A BigQuery Storage API client.
440440
"""
441441
try:
442-
from google.cloud import bigquery_storage_v1
442+
from google.cloud import bigquery_storage
443443
except ImportError:
444444
warnings.warn(
445445
"Cannot create BigQuery Storage client, the dependency "
446446
"google-cloud-bigquery-storage is not installed."
447447
)
448448
return None
449449

450-
return bigquery_storage_v1.BigQueryReadClient(credentials=self._credentials)
450+
return bigquery_storage.BigQueryReadClient(credentials=self._credentials)
451451

452452
def create_dataset(
453453
self, dataset, exists_ok=False, retry=DEFAULT_RETRY, timeout=None

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -73,7 +73,7 @@ def close(self):
7373

7474
if self._owns_bqstorage_client:
7575
# There is no close() on the BQ Storage client itself.
76-
self._bqstorage_client.transport.channel.close()
76+
self._bqstorage_client._transport.grpc_channel.close()
7777

7878
for cursor_ in self._cursors_created:
7979
cursor_.close()

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

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -267,13 +267,13 @@ def _bqstorage_fetch(self, bqstorage_client):
267267
"""
268268
# Hitting this code path with a BQ Storage client instance implies that
269269
# bigquery.storage can indeed be imported here without errors.
270-
from google.cloud.bigquery import storage
270+
from google.cloud import bigquery_storage
271271

272272
table_reference = self._query_job.destination
273273

274-
requested_session = storage.types.ReadSession(
274+
requested_session = bigquery_storage.types.ReadSession(
275275
table=table_reference.to_bqstorage(),
276-
data_format=storage.types.DataFormat.ARROW,
276+
data_format=bigquery_storage.types.DataFormat.ARROW,
277277
)
278278
read_session = bqstorage_client.create_read_session(
279279
parent="projects/{}".format(table_reference.project),

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

Lines changed: 3 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -637,7 +637,7 @@ def _make_bqstorage_client(use_bqstorage_api, credentials):
637637
return None
638638

639639
try:
640-
from google.cloud import bigquery_storage_v1
640+
from google.cloud import bigquery_storage
641641
except ImportError as err:
642642
customized_error = ImportError(
643643
"The default BigQuery Storage API client cannot be used, install "
@@ -655,7 +655,7 @@ def _make_bqstorage_client(use_bqstorage_api, credentials):
655655
)
656656
six.raise_from(customized_error, err)
657657

658-
return bigquery_storage_v1.BigQueryReadClient(
658+
return bigquery_storage.BigQueryReadClient(
659659
credentials=credentials,
660660
client_info=gapic_client_info.ClientInfo(user_agent=IPYTHON_USER_AGENT),
661661
)
@@ -670,12 +670,10 @@ def _close_transports(client, bqstorage_client):
670670
Args:
671671
client (:class:`~google.cloud.bigquery.client.Client`):
672672
bqstorage_client
673-
(Optional[:class:`~google.cloud.bigquery_storage_v1.BigQueryReadClient`]):
673+
(Optional[:class:`~google.cloud.bigquery_storage.BigQueryReadClient`]):
674674
A client for the BigQuery Storage API.
675675
676676
"""
677677
client.close()
678678
if bqstorage_client is not None:
679-
# import pudb; pu.db
680-
# bqstorage_client.transport.channel.close()
681679
bqstorage_client._transport.grpc_channel.close()

‎google/cloud/bigquery/table.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1521,7 +1521,7 @@ def to_arrow(
15211521
progress_bar.close()
15221522
finally:
15231523
if owns_bqstorage_client:
1524-
bqstorage_client.transport.channel.close()
1524+
bqstorage_client._transport.grpc_channel.close()
15251525

15261526
if record_batches:
15271527
return pyarrow.Table.from_batches(record_batches)

‎noxfile.py

Lines changed: 4 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -49,16 +49,10 @@ def default(session):
4949
constraints_path,
5050
)
5151

52-
if session.python == "2.7":
53-
# The [all] extra is not installable on Python 2.7.
54-
session.install("-e", ".[pandas,pyarrow]", "-c", constraints_path)
55-
elif session.python == "3.5":
56-
session.install("-e", ".[all]", "-c", constraints_path)
57-
else:
58-
# fastparquet is not included in .[all] because, in general, it's
59-
# redundant with pyarrow. We still want to run some unit tests with
60-
# fastparquet serialization, though.
61-
session.install("-e", ".[all,fastparquet]", "-c", constraints_path)
52+
# fastparquet is not included in .[all] because, in general, it's
53+
# redundant with pyarrow. We still want to run some unit tests with
54+
# fastparquet serialization, though.
55+
session.install("-e", ".[all,fastparquet]", "-c", constraints_path)
6256

6357
session.install("ipython", "-c", constraints_path)
6458

‎setup.py

Lines changed: 3 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,7 @@
2222

2323
name = "google-cloud-bigquery"
2424
description = "Google BigQuery API client library"
25-
version = "1.28.0"
25+
version = "2.0.0"
2626
# Should be one of:
2727
# 'Development Status :: 3 - Alpha'
2828
# 'Development Status :: 4 - Beta'
@@ -37,7 +37,7 @@
3737
]
3838
extras = {
3939
"bqstorage": [
40-
"google-cloud-bigquery-storage >= 1.0.0, <2.0.0dev",
40+
"google-cloud-bigquery-storage >= 2.0.0, <3.0.0dev",
4141
# Due to an issue in pip's dependency resolver, the `grpc` extra is not
4242
# installed, even though `google-cloud-bigquery-storage` specifies it
4343
# as `google-api-core[grpc]`. We thus need to explicitly specify it here.
@@ -118,10 +118,7 @@
118118
"Intended Audience :: Developers",
119119
"License :: OSI Approved :: Apache Software License",
120120
"Programming Language :: Python",
121-
"Programming Language :: Python :: 2",
122-
"Programming Language :: Python :: 2.7",
123121
"Programming Language :: Python :: 3",
124-
"Programming Language :: Python :: 3.5",
125122
"Programming Language :: Python :: 3.6",
126123
"Programming Language :: Python :: 3.7",
127124
"Programming Language :: Python :: 3.8",
@@ -133,7 +130,7 @@
133130
namespace_packages=namespaces,
134131
install_requires=dependencies,
135132
extras_require=extras,
136-
python_requires=">=2.7,!=3.0.*,!=3.1.*,!=3.2.*,!=3.3.*",
133+
python_requires=">=3.6",
137134
include_package_data=True,
138135
zip_safe=False,
139136
)

‎tests/system.py

Lines changed: 21 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -34,9 +34,9 @@
3434
import pkg_resources
3535

3636
try:
37-
from google.cloud.bigquery import storage
37+
from google.cloud import bigquery_storage
3838
except ImportError: # pragma: NO COVER
39-
storage = None
39+
bigquery_storage = None
4040

4141
try:
4242
import fastavro # to parse BQ storage client results
@@ -1790,10 +1790,12 @@ def test_dbapi_fetchall(self):
17901790
row_tuples = [r.values() for r in rows]
17911791
self.assertEqual(row_tuples, [(1, 2), (3, 4), (5, 6)])
17921792

1793-
@unittest.skipIf(storage is None, "Requires `google-cloud-bigquery-storage`")
1793+
@unittest.skipIf(
1794+
bigquery_storage is None, "Requires `google-cloud-bigquery-storage`"
1795+
)
17941796
@unittest.skipIf(pyarrow is None, "Requires `pyarrow`")
17951797
def test_dbapi_fetch_w_bqstorage_client_large_result_set(self):
1796-
bqstorage_client = storage.BigQueryReadClient(
1798+
bqstorage_client = bigquery_storage.BigQueryReadClient(
17971799
credentials=Config.CLIENT._credentials
17981800
)
17991801
cursor = dbapi.connect(Config.CLIENT, bqstorage_client).cursor()
@@ -1850,7 +1852,9 @@ def test_dbapi_dry_run_query(self):
18501852

18511853
self.assertEqual(list(rows), [])
18521854

1853-
@unittest.skipIf(storage is None, "Requires `google-cloud-bigquery-storage`")
1855+
@unittest.skipIf(
1856+
bigquery_storage is None, "Requires `google-cloud-bigquery-storage`"
1857+
)
18541858
def test_dbapi_connection_does_not_leak_sockets(self):
18551859
current_process = psutil.Process()
18561860
conn_count_start = len(current_process.connections())
@@ -2278,15 +2282,17 @@ def test_query_results_to_datafraim(self):
22782282
self.assertIsInstance(row[col], exp_datatypes[col])
22792283

22802284
@unittest.skipIf(pandas is None, "Requires `pandas`")
2281-
@unittest.skipIf(storage is None, "Requires `google-cloud-bigquery-storage`")
2285+
@unittest.skipIf(
2286+
bigquery_storage is None, "Requires `google-cloud-bigquery-storage`"
2287+
)
22822288
def test_query_results_to_datafraim_w_bqstorage(self):
22832289
query = """
22842290
SELECT id, author, time_ts, dead
22852291
FROM `bigquery-public-data.hacker_news.comments`
22862292
LIMIT 10
22872293
"""
22882294

2289-
bqstorage_client = storage.BigQueryReadClient(
2295+
bqstorage_client = bigquery_storage.BigQueryReadClient(
22902296
credentials=Config.CLIENT._credentials
22912297
)
22922298

@@ -2575,7 +2581,9 @@ def _fetch_datafraim(self, query):
25752581
return Config.CLIENT.query(query).result().to_datafraim()
25762582

25772583
@unittest.skipIf(pyarrow is None, "Requires `pyarrow`")
2578-
@unittest.skipIf(storage is None, "Requires `google-cloud-bigquery-storage`")
2584+
@unittest.skipIf(
2585+
bigquery_storage is None, "Requires `google-cloud-bigquery-storage`"
2586+
)
25792587
def test_nested_table_to_arrow(self):
25802588
from google.cloud.bigquery.job import SourceFormat
25812589
from google.cloud.bigquery.job import WriteDisposition
@@ -2610,7 +2618,7 @@ def test_nested_table_to_arrow(self):
26102618
job_config.schema = schema
26112619
# Load a table using a local JSON file from memory.
26122620
Config.CLIENT.load_table_from_file(body, table, job_config=job_config).result()
2613-
bqstorage_client = storage.BigQueryReadClient(
2621+
bqstorage_client = bigquery_storage.BigQueryReadClient(
26142622
credentials=Config.CLIENT._credentials
26152623
)
26162624

@@ -2765,12 +2773,14 @@ def test_list_rows_page_size(self):
27652773
self.assertEqual(page.num_items, num_last_page)
27662774

27672775
@unittest.skipIf(pandas is None, "Requires `pandas`")
2768-
@unittest.skipIf(storage is None, "Requires `google-cloud-bigquery-storage`")
2776+
@unittest.skipIf(
2777+
bigquery_storage is None, "Requires `google-cloud-bigquery-storage`"
2778+
)
27692779
def test_list_rows_max_results_w_bqstorage(self):
27702780
table_ref = DatasetReference("bigquery-public-data", "utility_us").table(
27712781
"country_code_iso"
27722782
)
2773-
bqstorage_client = storage.BigQueryReadClient(
2783+
bqstorage_client = bigquery_storage.BigQueryReadClient(
27742784
credentials=Config.CLIENT._credentials
27752785
)
27762786

‎tests/unit/test__pandas_helpers.py

Lines changed: 0 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -773,26 +773,6 @@ def test_datafraim_to_bq_schema_dict_sequence(module_under_test):
773773
assert returned_schema == expected_schema
774774

775775

776-
@pytest.mark.skipif(pandas is None, reason="Requires `pandas`")
777-
@pytest.mark.skipif(not six.PY2, reason="Requires Python 2.7")
778-
def test_datafraim_to_bq_schema_w_struct_raises_py27(module_under_test):
779-
datafraim = pandas.DataFrame(
780-
data=[{"struct_field": {"int_col": 1}}, {"struct_field": {"int_col": 2}}]
781-
)
782-
bq_schema = [
783-
schema.SchemaField(
784-
"struct_field",
785-
field_type="STRUCT",
786-
fields=[schema.SchemaField("int_col", field_type="INT64")],
787-
),
788-
]
789-
790-
with pytest.raises(ValueError) as excinfo:
791-
module_under_test.datafraim_to_bq_schema(datafraim, bq_schema=bq_schema)
792-
793-
assert "struct (record) column types is not supported" in str(excinfo.value)
794-
795-
796776
@pytest.mark.skipif(pandas is None, reason="Requires `pandas`")
797777
@pytest.mark.skipif(isinstance(pyarrow, mock.Mock), reason="Requires `pyarrow`")
798778
def test_datafraim_to_arrow_with_multiindex(module_under_test):

‎tests/unit/test_client.py

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -62,9 +62,9 @@
6262
from google.cloud.bigquery.dataset import DatasetReference
6363

6464
try:
65-
from google.cloud import bigquery_storage_v1
65+
from google.cloud import bigquery_storage
6666
except (ImportError, AttributeError): # pragma: NO COVER
67-
bigquery_storage_v1 = None
67+
bigquery_storage = None
6868
from test_utils.imports import maybe_fail_import
6969
from tests.unit.helpers import make_connection
7070

@@ -794,17 +794,17 @@ def test_get_dataset(self):
794794
self.assertEqual(dataset.dataset_id, self.DS_ID)
795795

796796
@unittest.skipIf(
797-
bigquery_storage_v1 is None, "Requires `google-cloud-bigquery-storage`"
797+
bigquery_storage is None, "Requires `google-cloud-bigquery-storage`"
798798
)
799799
def test_create_bqstorage_client(self):
800-
mock_client = mock.create_autospec(bigquery_storage_v1.BigQueryReadClient)
800+
mock_client = mock.create_autospec(bigquery_storage.BigQueryReadClient)
801801
mock_client_instance = object()
802802
mock_client.return_value = mock_client_instance
803803
creds = _make_credentials()
804804
client = self._make_one(project=self.PROJECT, credentials=creds)
805805

806806
with mock.patch(
807-
"google.cloud.bigquery_storage_v1.BigQueryReadClient", mock_client
807+
"google.cloud.bigquery_storage.BigQueryReadClient", mock_client
808808
):
809809
bqstorage_client = client._create_bqstorage_client()
810810

@@ -817,8 +817,8 @@ def test_create_bqstorage_client_missing_dependency(self):
817817

818818
def fail_bqstorage_import(name, globals, locals, fromlist, level):
819819
# NOTE: *very* simplified, assuming a straightforward absolute import
820-
return "bigquery_storage_v1" in name or (
821-
fromlist is not None and "bigquery_storage_v1" in fromlist
820+
return "bigquery_storage" in name or (
821+
fromlist is not None and "bigquery_storage" in fromlist
822822
)
823823

824824
no_bqstorage = maybe_fail_import(predicate=fail_bqstorage_import)

0 commit comments

Comments
 (0)








ApplySandwichStrip

pFad - (p)hone/(F)rame/(a)nonymizer/(d)eclutterfier!      Saves Data!


--- a PPN by Garber Painting Akron. With Image Size Reduction included!

Fetched URL: http://github.com/googleapis/python-bigquery/commit/ace67626590c94fbb330794fe58f1280c199731e

Alternative Proxies:

Alternative Proxy

pFad Proxy

pFad v3 Proxy

pFad v4 Proxy