diff --git a/google/cloud/bigquery/_pandas_helpers.py b/google/cloud/bigquery/_pandas_helpers.py index 5247a7d44..596527880 100644 --- a/google/cloud/bigquery/_pandas_helpers.py +++ b/google/cloud/bigquery/_pandas_helpers.py @@ -22,11 +22,6 @@ import six from six.moves import queue -try: - from google.cloud import bigquery_storage_v1 -except ImportError: # pragma: NO COVER - bigquery_storage_v1 = None - try: import pandas except ImportError: # pragma: NO COVER @@ -613,7 +608,7 @@ def _download_table_bqstorage( # Passing a BQ Storage client in implies that the BigQuery Storage library # is available and can be imported. - from google.cloud.bigquery import storage + from google.cloud import bigquery_storage if "$" in table.table_id: raise ValueError( @@ -624,8 +619,8 @@ def _download_table_bqstorage( requested_streams = 1 if preserve_order else 0 - requested_session = storage.types.ReadSession( - table=table.to_bqstorage(), data_format=storage.types.DataFormat.ARROW + requested_session = bigquery_storage.types.ReadSession( + table=table.to_bqstorage(), data_format=bigquery_storage.types.DataFormat.ARROW ) if selected_fields is not None: for field in selected_fields: diff --git a/google/cloud/bigquery/client.py b/google/cloud/bigquery/client.py index d2aa45999..942280ff7 100644 --- a/google/cloud/bigquery/client.py +++ b/google/cloud/bigquery/client.py @@ -435,11 +435,11 @@ def _create_bqstorage_client(self): warning and return ``None``. Returns: - Optional[google.cloud.bigquery_storage_v1.BigQueryReadClient]: + Optional[google.cloud.bigquery_storage.BigQueryReadClient]: A BigQuery Storage API client. """ try: - from google.cloud import bigquery_storage_v1 + from google.cloud import bigquery_storage except ImportError: warnings.warn( "Cannot create BigQuery Storage client, the dependency " @@ -447,7 +447,7 @@ def _create_bqstorage_client(self): ) return None - return bigquery_storage_v1.BigQueryReadClient(credentials=self._credentials) + return bigquery_storage.BigQueryReadClient(credentials=self._credentials) def create_dataset( self, dataset, exists_ok=False, retry=DEFAULT_RETRY, timeout=None diff --git a/google/cloud/bigquery/dbapi/connection.py b/google/cloud/bigquery/dbapi/connection.py index 464b0fd06..300c77dc9 100644 --- a/google/cloud/bigquery/dbapi/connection.py +++ b/google/cloud/bigquery/dbapi/connection.py @@ -73,7 +73,7 @@ def close(self): if self._owns_bqstorage_client: # There is no close() on the BQ Storage client itself. - self._bqstorage_client.transport.channel.close() + self._bqstorage_client._transport.grpc_channel.close() for cursor_ in self._cursors_created: cursor_.close() diff --git a/google/cloud/bigquery/dbapi/cursor.py b/google/cloud/bigquery/dbapi/cursor.py index 32e8b1cee..b9f404335 100644 --- a/google/cloud/bigquery/dbapi/cursor.py +++ b/google/cloud/bigquery/dbapi/cursor.py @@ -267,13 +267,13 @@ def _bqstorage_fetch(self, bqstorage_client): """ # Hitting this code path with a BQ Storage client instance implies that # bigquery.storage can indeed be imported here without errors. - from google.cloud.bigquery import storage + from google.cloud import bigquery_storage table_reference = self._query_job.destination - requested_session = storage.types.ReadSession( + requested_session = bigquery_storage.types.ReadSession( table=table_reference.to_bqstorage(), - data_format=storage.types.DataFormat.ARROW, + data_format=bigquery_storage.types.DataFormat.ARROW, ) read_session = bqstorage_client.create_read_session( parent="projects/{}".format(table_reference.project), diff --git a/google/cloud/bigquery/magics/magics.py b/google/cloud/bigquery/magics/magics.py index 9b7874279..22175ee45 100644 --- a/google/cloud/bigquery/magics/magics.py +++ b/google/cloud/bigquery/magics/magics.py @@ -637,7 +637,7 @@ def _make_bqstorage_client(use_bqstorage_api, credentials): return None try: - from google.cloud import bigquery_storage_v1 + from google.cloud import bigquery_storage except ImportError as err: customized_error = ImportError( "The default BigQuery Storage API client cannot be used, install " @@ -655,7 +655,7 @@ def _make_bqstorage_client(use_bqstorage_api, credentials): ) six.raise_from(customized_error, err) - return bigquery_storage_v1.BigQueryReadClient( + return bigquery_storage.BigQueryReadClient( credentials=credentials, client_info=gapic_client_info.ClientInfo(user_agent=IPYTHON_USER_AGENT), ) @@ -670,12 +670,10 @@ def _close_transports(client, bqstorage_client): Args: client (:class:`~google.cloud.bigquery.client.Client`): bqstorage_client - (Optional[:class:`~google.cloud.bigquery_storage_v1.BigQueryReadClient`]): + (Optional[:class:`~google.cloud.bigquery_storage.BigQueryReadClient`]): A client for the BigQuery Storage API. """ client.close() if bqstorage_client is not None: - # import pudb; pu.db - # bqstorage_client.transport.channel.close() bqstorage_client._transport.grpc_channel.close() diff --git a/google/cloud/bigquery/table.py b/google/cloud/bigquery/table.py index d42b56e4b..f02a2f46c 100644 --- a/google/cloud/bigquery/table.py +++ b/google/cloud/bigquery/table.py @@ -1521,7 +1521,7 @@ def to_arrow( progress_bar.close() finally: if owns_bqstorage_client: - bqstorage_client.transport.channel.close() + bqstorage_client._transport.grpc_channel.close() if record_batches: return pyarrow.Table.from_batches(record_batches) diff --git a/noxfile.py b/noxfile.py index 90f023add..308c0adb7 100644 --- a/noxfile.py +++ b/noxfile.py @@ -49,16 +49,10 @@ def default(session): constraints_path, ) - if session.python == "2.7": - # The [all] extra is not installable on Python 2.7. - session.install("-e", ".[pandas,pyarrow]", "-c", constraints_path) - elif session.python == "3.5": - session.install("-e", ".[all]", "-c", constraints_path) - else: - # fastparquet is not included in .[all] because, in general, it's - # redundant with pyarrow. We still want to run some unit tests with - # fastparquet serialization, though. - session.install("-e", ".[all,fastparquet]", "-c", constraints_path) + # fastparquet is not included in .[all] because, in general, it's + # redundant with pyarrow. We still want to run some unit tests with + # fastparquet serialization, though. + session.install("-e", ".[all,fastparquet]", "-c", constraints_path) session.install("ipython", "-c", constraints_path) diff --git a/setup.py b/setup.py index 73d9a03ca..eb86bd812 100644 --- a/setup.py +++ b/setup.py @@ -22,7 +22,7 @@ name = "google-cloud-bigquery" description = "Google BigQuery API client library" -version = "1.28.0" +version = "2.0.0" # Should be one of: # 'Development Status :: 3 - Alpha' # 'Development Status :: 4 - Beta' @@ -37,7 +37,7 @@ ] extras = { "bqstorage": [ - "google-cloud-bigquery-storage >= 1.0.0, <2.0.0dev", + "google-cloud-bigquery-storage >= 2.0.0, <3.0.0dev", # Due to an issue in pip's dependency resolver, the `grpc` extra is not # installed, even though `google-cloud-bigquery-storage` specifies it # as `google-api-core[grpc]`. We thus need to explicitly specify it here. @@ -118,10 +118,7 @@ "Intended Audience :: Developers", "License :: OSI Approved :: Apache Software License", "Programming Language :: Python", - "Programming Language :: Python :: 2", - "Programming Language :: Python :: 2.7", "Programming Language :: Python :: 3", - "Programming Language :: Python :: 3.5", "Programming Language :: Python :: 3.6", "Programming Language :: Python :: 3.7", "Programming Language :: Python :: 3.8", @@ -133,7 +130,7 @@ namespace_packages=namespaces, install_requires=dependencies, extras_require=extras, - python_requires=">=2.7,!=3.0.*,!=3.1.*,!=3.2.*,!=3.3.*", + python_requires=">=3.6", include_package_data=True, zip_safe=False, ) diff --git a/tests/system.py b/tests/system.py index d8c45840c..f6e3a94ba 100644 --- a/tests/system.py +++ b/tests/system.py @@ -34,9 +34,9 @@ import pkg_resources try: - from google.cloud.bigquery import storage + from google.cloud import bigquery_storage except ImportError: # pragma: NO COVER - storage = None + bigquery_storage = None try: import fastavro # to parse BQ storage client results @@ -1790,10 +1790,12 @@ def test_dbapi_fetchall(self): row_tuples = [r.values() for r in rows] self.assertEqual(row_tuples, [(1, 2), (3, 4), (5, 6)]) - @unittest.skipIf(storage is None, "Requires `google-cloud-bigquery-storage`") + @unittest.skipIf( + bigquery_storage is None, "Requires `google-cloud-bigquery-storage`" + ) @unittest.skipIf(pyarrow is None, "Requires `pyarrow`") def test_dbapi_fetch_w_bqstorage_client_large_result_set(self): - bqstorage_client = storage.BigQueryReadClient( + bqstorage_client = bigquery_storage.BigQueryReadClient( credentials=Config.CLIENT._credentials ) cursor = dbapi.connect(Config.CLIENT, bqstorage_client).cursor() @@ -1850,7 +1852,9 @@ def test_dbapi_dry_run_query(self): self.assertEqual(list(rows), []) - @unittest.skipIf(storage is None, "Requires `google-cloud-bigquery-storage`") + @unittest.skipIf( + bigquery_storage is None, "Requires `google-cloud-bigquery-storage`" + ) def test_dbapi_connection_does_not_leak_sockets(self): current_process = psutil.Process() conn_count_start = len(current_process.connections()) @@ -2278,7 +2282,9 @@ def test_query_results_to_dataframe(self): self.assertIsInstance(row[col], exp_datatypes[col]) @unittest.skipIf(pandas is None, "Requires `pandas`") - @unittest.skipIf(storage is None, "Requires `google-cloud-bigquery-storage`") + @unittest.skipIf( + bigquery_storage is None, "Requires `google-cloud-bigquery-storage`" + ) def test_query_results_to_dataframe_w_bqstorage(self): query = """ SELECT id, author, time_ts, dead @@ -2286,7 +2292,7 @@ def test_query_results_to_dataframe_w_bqstorage(self): LIMIT 10 """ - bqstorage_client = storage.BigQueryReadClient( + bqstorage_client = bigquery_storage.BigQueryReadClient( credentials=Config.CLIENT._credentials ) @@ -2575,7 +2581,9 @@ def _fetch_dataframe(self, query): return Config.CLIENT.query(query).result().to_dataframe() @unittest.skipIf(pyarrow is None, "Requires `pyarrow`") - @unittest.skipIf(storage is None, "Requires `google-cloud-bigquery-storage`") + @unittest.skipIf( + bigquery_storage is None, "Requires `google-cloud-bigquery-storage`" + ) def test_nested_table_to_arrow(self): from google.cloud.bigquery.job import SourceFormat from google.cloud.bigquery.job import WriteDisposition @@ -2610,7 +2618,7 @@ def test_nested_table_to_arrow(self): job_config.schema = schema # Load a table using a local JSON file from memory. Config.CLIENT.load_table_from_file(body, table, job_config=job_config).result() - bqstorage_client = storage.BigQueryReadClient( + bqstorage_client = bigquery_storage.BigQueryReadClient( credentials=Config.CLIENT._credentials ) @@ -2765,12 +2773,14 @@ def test_list_rows_page_size(self): self.assertEqual(page.num_items, num_last_page) @unittest.skipIf(pandas is None, "Requires `pandas`") - @unittest.skipIf(storage is None, "Requires `google-cloud-bigquery-storage`") + @unittest.skipIf( + bigquery_storage is None, "Requires `google-cloud-bigquery-storage`" + ) def test_list_rows_max_results_w_bqstorage(self): table_ref = DatasetReference("bigquery-public-data", "utility_us").table( "country_code_iso" ) - bqstorage_client = storage.BigQueryReadClient( + bqstorage_client = bigquery_storage.BigQueryReadClient( credentials=Config.CLIENT._credentials ) diff --git a/tests/unit/test__pandas_helpers.py b/tests/unit/test__pandas_helpers.py index f4355072a..e229e04a2 100644 --- a/tests/unit/test__pandas_helpers.py +++ b/tests/unit/test__pandas_helpers.py @@ -773,26 +773,6 @@ def test_dataframe_to_bq_schema_dict_sequence(module_under_test): assert returned_schema == expected_schema -@pytest.mark.skipif(pandas is None, reason="Requires `pandas`") -@pytest.mark.skipif(not six.PY2, reason="Requires Python 2.7") -def test_dataframe_to_bq_schema_w_struct_raises_py27(module_under_test): - dataframe = pandas.DataFrame( - data=[{"struct_field": {"int_col": 1}}, {"struct_field": {"int_col": 2}}] - ) - bq_schema = [ - schema.SchemaField( - "struct_field", - field_type="STRUCT", - fields=[schema.SchemaField("int_col", field_type="INT64")], - ), - ] - - with pytest.raises(ValueError) as excinfo: - module_under_test.dataframe_to_bq_schema(dataframe, bq_schema=bq_schema) - - assert "struct (record) column types is not supported" in str(excinfo.value) - - @pytest.mark.skipif(pandas is None, reason="Requires `pandas`") @pytest.mark.skipif(isinstance(pyarrow, mock.Mock), reason="Requires `pyarrow`") def test_dataframe_to_arrow_with_multiindex(module_under_test): diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py index c4c604ed0..a5c259b3a 100644 --- a/tests/unit/test_client.py +++ b/tests/unit/test_client.py @@ -62,9 +62,9 @@ from google.cloud.bigquery.dataset import DatasetReference try: - from google.cloud import bigquery_storage_v1 + from google.cloud import bigquery_storage except (ImportError, AttributeError): # pragma: NO COVER - bigquery_storage_v1 = None + bigquery_storage = None from test_utils.imports import maybe_fail_import from tests.unit.helpers import make_connection @@ -794,17 +794,17 @@ def test_get_dataset(self): self.assertEqual(dataset.dataset_id, self.DS_ID) @unittest.skipIf( - bigquery_storage_v1 is None, "Requires `google-cloud-bigquery-storage`" + bigquery_storage is None, "Requires `google-cloud-bigquery-storage`" ) def test_create_bqstorage_client(self): - mock_client = mock.create_autospec(bigquery_storage_v1.BigQueryReadClient) + mock_client = mock.create_autospec(bigquery_storage.BigQueryReadClient) mock_client_instance = object() mock_client.return_value = mock_client_instance creds = _make_credentials() client = self._make_one(project=self.PROJECT, credentials=creds) with mock.patch( - "google.cloud.bigquery_storage_v1.BigQueryReadClient", mock_client + "google.cloud.bigquery_storage.BigQueryReadClient", mock_client ): bqstorage_client = client._create_bqstorage_client() @@ -817,8 +817,8 @@ def test_create_bqstorage_client_missing_dependency(self): def fail_bqstorage_import(name, globals, locals, fromlist, level): # NOTE: *very* simplified, assuming a straightforward absolute import - return "bigquery_storage_v1" in name or ( - fromlist is not None and "bigquery_storage_v1" in fromlist + return "bigquery_storage" in name or ( + fromlist is not None and "bigquery_storage" in fromlist ) no_bqstorage = maybe_fail_import(predicate=fail_bqstorage_import) diff --git a/tests/unit/test_dbapi_connection.py b/tests/unit/test_dbapi_connection.py index 0f1be45ee..b59b7e70f 100644 --- a/tests/unit/test_dbapi_connection.py +++ b/tests/unit/test_dbapi_connection.py @@ -19,9 +19,9 @@ import six try: - from google.cloud import bigquery_storage_v1 + from google.cloud import bigquery_storage except ImportError: # pragma: NO COVER - bigquery_storage_v1 = None + bigquery_storage = None class TestConnection(unittest.TestCase): @@ -41,13 +41,11 @@ def _mock_client(self): return mock_client def _mock_bqstorage_client(self): - if bigquery_storage_v1 is None: + if bigquery_storage is None: return None - mock_client = mock.create_autospec( - bigquery_storage_v1.client.BigQueryReadClient - ) - mock_client.transport = mock.Mock(spec=["channel"]) - mock_client.transport.channel = mock.Mock(spec=["close"]) + mock_client = mock.create_autospec(bigquery_storage.BigQueryReadClient) + mock_client._transport = mock.Mock(spec=["channel"]) + mock_client._transport.grpc_channel = mock.Mock(spec=["close"]) return mock_client def test_ctor_wo_bqstorage_client(self): @@ -63,7 +61,7 @@ def test_ctor_wo_bqstorage_client(self): self.assertIs(connection._bqstorage_client, mock_bqstorage_client) @unittest.skipIf( - bigquery_storage_v1 is None, "Requires `google-cloud-bigquery-storage`" + bigquery_storage is None, "Requires `google-cloud-bigquery-storage`" ) def test_ctor_w_bqstorage_client(self): from google.cloud.bigquery.dbapi import Connection @@ -101,7 +99,7 @@ def test_connect_w_client(self): self.assertIs(connection._bqstorage_client, mock_bqstorage_client) @unittest.skipIf( - bigquery_storage_v1 is None, "Requires `google-cloud-bigquery-storage`" + bigquery_storage is None, "Requires `google-cloud-bigquery-storage`" ) def test_connect_w_both_clients(self): from google.cloud.bigquery.dbapi import connect @@ -130,7 +128,7 @@ def test_raises_error_if_closed(self): getattr(connection, method)() @unittest.skipIf( - bigquery_storage_v1 is None, "Requires `google-cloud-bigquery-storage`" + bigquery_storage is None, "Requires `google-cloud-bigquery-storage`" ) def test_close_closes_all_created_bigquery_clients(self): client = self._mock_client() @@ -150,10 +148,10 @@ def test_close_closes_all_created_bigquery_clients(self): connection.close() self.assertTrue(client.close.called) - self.assertTrue(bqstorage_client.transport.channel.close.called) + self.assertTrue(bqstorage_client._transport.grpc_channel.close.called) @unittest.skipIf( - bigquery_storage_v1 is None, "Requires `google-cloud-bigquery-storage`" + bigquery_storage is None, "Requires `google-cloud-bigquery-storage`" ) def test_close_does_not_close_bigquery_clients_passed_to_it(self): client = self._mock_client() @@ -163,7 +161,7 @@ def test_close_does_not_close_bigquery_clients_passed_to_it(self): connection.close() self.assertFalse(client.close.called) - self.assertFalse(bqstorage_client.transport.channel.called) + self.assertFalse(bqstorage_client._transport.grpc_channel.close.called) def test_close_closes_all_created_cursors(self): connection = self._make_one(client=self._mock_client()) diff --git a/tests/unit/test_dbapi_cursor.py b/tests/unit/test_dbapi_cursor.py index ec05e5d47..9a1a6b1e8 100644 --- a/tests/unit/test_dbapi_cursor.py +++ b/tests/unit/test_dbapi_cursor.py @@ -26,9 +26,9 @@ from google.api_core import exceptions try: - from google.cloud.bigquery import storage + from google.cloud import bigquery_storage except ImportError: # pragma: NO COVER - storage = None + bigquery_storage = None from tests.unit.helpers import _to_pyarrow @@ -79,10 +79,10 @@ def _mock_bqstorage_client(self, rows=None, stream_count=0): if rows is None: rows = [] - mock_client = mock.create_autospec(storage.BigQueryReadClient) + mock_client = mock.create_autospec(bigquery_storage.BigQueryReadClient) mock_read_session = mock.MagicMock( streams=[ - storage.types.ReadStream(name="streams/stream_{}".format(i)) + bigquery_storage.types.ReadStream(name="streams/stream_{}".format(i)) for i in range(stream_count) ] ) @@ -272,7 +272,9 @@ def test_fetchall_w_row(self): self.assertEqual(len(rows), 1) self.assertEqual(rows[0], (1,)) - @unittest.skipIf(storage is None, "Requires `google-cloud-bigquery-storage`") + @unittest.skipIf( + bigquery_storage is None, "Requires `google-cloud-bigquery-storage`" + ) @unittest.skipIf(pyarrow is None, "Requires `pyarrow`") def test_fetchall_w_bqstorage_client_fetch_success(self): from google.cloud.bigquery import dbapi @@ -324,7 +326,9 @@ def test_fetchall_w_bqstorage_client_fetch_success(self): self.assertEqual(sorted_row_data, expected_row_data) - @unittest.skipIf(storage is None, "Requires `google-cloud-bigquery-storage`") + @unittest.skipIf( + bigquery_storage is None, "Requires `google-cloud-bigquery-storage`" + ) def test_fetchall_w_bqstorage_client_fetch_no_rows(self): from google.cloud.bigquery import dbapi @@ -345,7 +349,9 @@ def test_fetchall_w_bqstorage_client_fetch_no_rows(self): # check the data returned self.assertEqual(rows, []) - @unittest.skipIf(storage is None, "Requires `google-cloud-bigquery-storage`") + @unittest.skipIf( + bigquery_storage is None, "Requires `google-cloud-bigquery-storage`" + ) def test_fetchall_w_bqstorage_client_fetch_error_no_fallback(self): from google.cloud.bigquery import dbapi from google.cloud.bigquery import table diff --git a/tests/unit/test_job.py b/tests/unit/test_job.py index fb6a46bd6..fb042e18c 100644 --- a/tests/unit/test_job.py +++ b/tests/unit/test_job.py @@ -35,9 +35,9 @@ except ImportError: # pragma: NO COVER pyarrow = None try: - from google.cloud import bigquery_storage_v1 + from google.cloud import bigquery_storage except (ImportError, AttributeError): # pragma: NO COVER - bigquery_storage_v1 = None + bigquery_storage = None try: from tqdm import tqdm except (ImportError, AttributeError): # pragma: NO COVER @@ -5667,7 +5667,7 @@ def test_to_dataframe_ddl_query(self): @unittest.skipIf(pandas is None, "Requires `pandas`") @unittest.skipIf( - bigquery_storage_v1 is None, "Requires `google-cloud-bigquery-storage`" + bigquery_storage is None, "Requires `google-cloud-bigquery-storage`" ) def test_to_dataframe_bqstorage(self): query_resource = { @@ -5685,8 +5685,8 @@ def test_to_dataframe_bqstorage(self): client = _make_client(self.PROJECT, connection=connection) resource = self._make_resource(ended=True) job = self._get_target_class().from_api_repr(resource, client) - bqstorage_client = mock.create_autospec(bigquery_storage_v1.BigQueryReadClient) - session = bigquery_storage_v1.types.ReadSession() + bqstorage_client = mock.create_autospec(bigquery_storage.BigQueryReadClient) + session = bigquery_storage.types.ReadSession() session.avro_schema.schema = json.dumps( { "type": "record", @@ -5704,9 +5704,9 @@ def test_to_dataframe_bqstorage(self): destination_table = "projects/{projectId}/datasets/{datasetId}/tables/{tableId}".format( **resource["configuration"]["query"]["destinationTable"] ) - expected_session = bigquery_storage_v1.types.ReadSession( + expected_session = bigquery_storage.types.ReadSession( table=destination_table, - data_format=bigquery_storage_v1.enums.DataFormat.ARROW, + data_format=bigquery_storage.types.DataFormat.ARROW, ) bqstorage_client.create_read_session.assert_called_once_with( parent="projects/{}".format(self.PROJECT), @@ -6259,7 +6259,7 @@ def test__contains_order_by(query, expected): @pytest.mark.skipif(pandas is None, reason="Requires `pandas`") @pytest.mark.skipif( - bigquery_storage_v1 is None, reason="Requires `google-cloud-bigquery-storage`" + bigquery_storage is None, reason="Requires `google-cloud-bigquery-storage`" ) @pytest.mark.parametrize( "query", @@ -6295,8 +6295,8 @@ def test_to_dataframe_bqstorage_preserve_order(query): connection = _make_connection(get_query_results_resource, job_resource) client = _make_client(connection=connection) job = target_class.from_api_repr(job_resource, client) - bqstorage_client = mock.create_autospec(bigquery_storage_v1.BigQueryReadClient) - session = bigquery_storage_v1.types.ReadSession() + bqstorage_client = mock.create_autospec(bigquery_storage.BigQueryReadClient) + session = bigquery_storage.types.ReadSession() session.avro_schema.schema = json.dumps( { "type": "record", @@ -6314,8 +6314,8 @@ def test_to_dataframe_bqstorage_preserve_order(query): destination_table = "projects/{projectId}/datasets/{datasetId}/tables/{tableId}".format( **job_resource["configuration"]["query"]["destinationTable"] ) - expected_session = bigquery_storage_v1.types.ReadSession( - table=destination_table, data_format=bigquery_storage_v1.enums.DataFormat.ARROW, + expected_session = bigquery_storage.types.ReadSession( + table=destination_table, data_format=bigquery_storage.types.DataFormat.ARROW, ) bqstorage_client.create_read_session.assert_called_once_with( parent="projects/test-project", diff --git a/tests/unit/test_magics.py b/tests/unit/test_magics.py index c4527c837..20be6b755 100644 --- a/tests/unit/test_magics.py +++ b/tests/unit/test_magics.py @@ -41,7 +41,7 @@ io = pytest.importorskip("IPython.utils.io") tools = pytest.importorskip("IPython.testing.tools") interactiveshell = pytest.importorskip("IPython.terminal.interactiveshell") -bigquery_storage_v1 = pytest.importorskip("google.cloud.bigquery_storage_v1") +bigquery_storage = pytest.importorskip("google.cloud.bigquery_storage") @pytest.fixture(scope="session") @@ -83,8 +83,8 @@ def missing_bq_storage(): def fail_if(name, globals, locals, fromlist, level): # NOTE: *very* simplified, assuming a straightforward absolute import - return "bigquery_storage_v1" in name or ( - fromlist is not None and "bigquery_storage_v1" in fromlist + return "bigquery_storage" in name or ( + fromlist is not None and "bigquery_storage" in fromlist ) return maybe_fail_import(predicate=fail_if) @@ -314,14 +314,14 @@ def test__make_bqstorage_client_false(): @pytest.mark.skipif( - bigquery_storage_v1 is None, reason="Requires `google-cloud-bigquery-storage`" + bigquery_storage is None, reason="Requires `google-cloud-bigquery-storage`" ) def test__make_bqstorage_client_true(): credentials_mock = mock.create_autospec( google.auth.credentials.Credentials, instance=True ) got = magics._make_bqstorage_client(True, credentials_mock) - assert isinstance(got, bigquery_storage_v1.BigQueryReadClient) + assert isinstance(got, bigquery_storage.BigQueryReadClient) def test__make_bqstorage_client_true_raises_import_error(missing_bq_storage): @@ -338,7 +338,7 @@ def test__make_bqstorage_client_true_raises_import_error(missing_bq_storage): @pytest.mark.skipif( - bigquery_storage_v1 is None, reason="Requires `google-cloud-bigquery-storage`" + bigquery_storage is None, reason="Requires `google-cloud-bigquery-storage`" ) @pytest.mark.skipif(pandas is None, reason="Requires `pandas`") def test__make_bqstorage_client_true_missing_gapic(missing_grpcio_lib): @@ -396,7 +396,7 @@ def test_extension_load(): @pytest.mark.usefixtures("ipython_interactive") @pytest.mark.skipif(pandas is None, reason="Requires `pandas`") @pytest.mark.skipif( - bigquery_storage_v1 is None, reason="Requires `google-cloud-bigquery-storage`" + bigquery_storage is None, reason="Requires `google-cloud-bigquery-storage`" ) def test_bigquery_magic_without_optional_arguments(monkeypatch): ip = IPython.get_ipython() @@ -410,14 +410,14 @@ def test_bigquery_magic_without_optional_arguments(monkeypatch): monkeypatch.setattr(magics.context, "_credentials", mock_credentials) # Mock out the BigQuery Storage API. - bqstorage_mock = mock.create_autospec(bigquery_storage_v1.BigQueryReadClient) + bqstorage_mock = mock.create_autospec(bigquery_storage.BigQueryReadClient) bqstorage_instance_mock = mock.create_autospec( - bigquery_storage_v1.BigQueryReadClient, instance=True + bigquery_storage.BigQueryReadClient, instance=True ) - bqstorage_instance_mock.transport = mock.Mock() + bqstorage_instance_mock._transport = mock.Mock() bqstorage_mock.return_value = bqstorage_instance_mock bqstorage_client_patch = mock.patch( - "google.cloud.bigquery_storage_v1.BigQueryReadClient", bqstorage_mock + "google.cloud.bigquery_storage.BigQueryReadClient", bqstorage_mock ) sql = "SELECT 17 AS num" @@ -559,7 +559,7 @@ def test_bigquery_magic_clears_display_in_verbose_mode(): @pytest.mark.usefixtures("ipython_interactive") @pytest.mark.skipif( - bigquery_storage_v1 is None, reason="Requires `google-cloud-bigquery-storage`" + bigquery_storage is None, reason="Requires `google-cloud-bigquery-storage`" ) def test_bigquery_magic_with_bqstorage_from_argument(monkeypatch): ip = IPython.get_ipython() @@ -573,14 +573,14 @@ def test_bigquery_magic_with_bqstorage_from_argument(monkeypatch): monkeypatch.setattr(magics.context, "_credentials", mock_credentials) # Mock out the BigQuery Storage API. - bqstorage_mock = mock.create_autospec(bigquery_storage_v1.BigQueryReadClient) + bqstorage_mock = mock.create_autospec(bigquery_storage.BigQueryReadClient) bqstorage_instance_mock = mock.create_autospec( - bigquery_storage_v1.BigQueryReadClient, instance=True + bigquery_storage.BigQueryReadClient, instance=True ) - bqstorage_instance_mock.transport = mock.Mock() + bqstorage_instance_mock._transport = mock.Mock() bqstorage_mock.return_value = bqstorage_instance_mock bqstorage_client_patch = mock.patch( - "google.cloud.bigquery_storage_v1.BigQueryReadClient", bqstorage_mock + "google.cloud.bigquery_storage.BigQueryReadClient", bqstorage_mock ) sql = "SELECT 17 AS num" @@ -623,7 +623,7 @@ def warning_match(warning): @pytest.mark.usefixtures("ipython_interactive") @pytest.mark.skipif( - bigquery_storage_v1 is None, reason="Requires `google-cloud-bigquery-storage`" + bigquery_storage is None, reason="Requires `google-cloud-bigquery-storage`" ) def test_bigquery_magic_with_rest_client_requested(monkeypatch): ip = IPython.get_ipython() @@ -637,9 +637,9 @@ def test_bigquery_magic_with_rest_client_requested(monkeypatch): monkeypatch.setattr(magics.context, "_credentials", mock_credentials) # Mock out the BigQuery Storage API. - bqstorage_mock = mock.create_autospec(bigquery_storage_v1.BigQueryReadClient) + bqstorage_mock = mock.create_autospec(bigquery_storage.BigQueryReadClient) bqstorage_client_patch = mock.patch( - "google.cloud.bigquery_storage_v1.BigQueryReadClient", bqstorage_mock + "google.cloud.bigquery_storage.BigQueryReadClient", bqstorage_mock ) sql = "SELECT 17 AS num" @@ -841,7 +841,7 @@ def test_bigquery_magic_w_table_id_and_destination_var(ipython_ns_cleanup): @pytest.mark.usefixtures("ipython_interactive") @pytest.mark.skipif( - bigquery_storage_v1 is None, reason="Requires `google-cloud-bigquery-storage`" + bigquery_storage is None, reason="Requires `google-cloud-bigquery-storage`" ) @pytest.mark.skipif(pandas is None, reason="Requires `pandas`") def test_bigquery_magic_w_table_id_and_bqstorage_client(): @@ -864,14 +864,14 @@ def test_bigquery_magic_w_table_id_and_bqstorage_client(): "google.cloud.bigquery.magics.magics.bigquery.Client", autospec=True ) - bqstorage_mock = mock.create_autospec(bigquery_storage_v1.BigQueryReadClient) + bqstorage_mock = mock.create_autospec(bigquery_storage.BigQueryReadClient) bqstorage_instance_mock = mock.create_autospec( - bigquery_storage_v1.BigQueryReadClient, instance=True + bigquery_storage.BigQueryReadClient, instance=True ) - bqstorage_instance_mock.transport = mock.Mock() + bqstorage_instance_mock._transport = mock.Mock() bqstorage_mock.return_value = bqstorage_instance_mock bqstorage_client_patch = mock.patch( - "google.cloud.bigquery_storage_v1.BigQueryReadClient", bqstorage_mock + "google.cloud.bigquery_storage.BigQueryReadClient", bqstorage_mock ) table_id = "bigquery-public-data.samples.shakespeare" diff --git a/tests/unit/test_table.py b/tests/unit/test_table.py index 5e5a1593a..12169658e 100644 --- a/tests/unit/test_table.py +++ b/tests/unit/test_table.py @@ -25,12 +25,12 @@ import google.api_core.exceptions try: - from google.cloud.bigquery import storage - from google.cloud.bigquery.storage_v1.services.big_query_read.transports import ( - grpc as big_query_read_grpc_transport + from google.cloud import bigquery_storage + from google.cloud.bigquery_storage_v1.services.big_query_read.transports import ( + grpc as big_query_read_grpc_transport, ) except ImportError: # pragma: NO COVER - storage = None + bigquery_storage = None big_query_read_grpc_transport = None try: @@ -1838,7 +1838,9 @@ def test_to_arrow_w_empty_table(self): self.assertEqual(child_field.type.value_type[1].name, "age") @unittest.skipIf(pyarrow is None, "Requires `pyarrow`") - @unittest.skipIf(storage is None, "Requires `google-cloud-bigquery-storage`") + @unittest.skipIf( + bigquery_storage is None, "Requires `google-cloud-bigquery-storage`" + ) def test_to_arrow_max_results_w_create_bqstorage_warning(self): from google.cloud.bigquery.schema import SchemaField @@ -1876,14 +1878,16 @@ def test_to_arrow_max_results_w_create_bqstorage_warning(self): mock_client._create_bqstorage_client.assert_not_called() @unittest.skipIf(pyarrow is None, "Requires `pyarrow`") - @unittest.skipIf(storage is None, "Requires `google-cloud-bigquery-storage`") + @unittest.skipIf( + bigquery_storage is None, "Requires `google-cloud-bigquery-storage`" + ) def test_to_arrow_w_bqstorage(self): from google.cloud.bigquery import schema from google.cloud.bigquery import table as mut from google.cloud.bigquery_storage_v1 import reader - bqstorage_client = mock.create_autospec(storage.BigQueryReadClient) - bqstorage_client.transport = mock.create_autospec( + bqstorage_client = mock.create_autospec(bigquery_storage.BigQueryReadClient) + bqstorage_client._transport = mock.create_autospec( big_query_read_grpc_transport.BigQueryReadGrpcTransport ) streams = [ @@ -1891,7 +1895,7 @@ def test_to_arrow_w_bqstorage(self): {"name": "/projects/proj/dataset/dset/tables/tbl/streams/1234"}, {"name": "/projects/proj/dataset/dset/tables/tbl/streams/5678"}, ] - session = storage.types.ReadSession(streams=streams) + session = bigquery_storage.types.ReadSession(streams=streams) arrow_schema = pyarrow.schema( [ pyarrow.field("colA", pyarrow.int64()), @@ -1952,21 +1956,23 @@ def test_to_arrow_w_bqstorage(self): self.assertEqual(actual_tbl.num_rows, total_rows) # Don't close the client if it was passed in. - bqstorage_client.transport.channel.close.assert_not_called() + bqstorage_client._transport.grpc_channel.close.assert_not_called() @unittest.skipIf(pyarrow is None, "Requires `pyarrow`") - @unittest.skipIf(storage is None, "Requires `google-cloud-bigquery-storage`") + @unittest.skipIf( + bigquery_storage is None, "Requires `google-cloud-bigquery-storage`" + ) def test_to_arrow_w_bqstorage_creates_client(self): from google.cloud.bigquery import schema from google.cloud.bigquery import table as mut mock_client = _mock_client() - bqstorage_client = mock.create_autospec(storage.BigQueryReadClient) - bqstorage_client.transport = mock.create_autospec( + bqstorage_client = mock.create_autospec(bigquery_storage.BigQueryReadClient) + bqstorage_client._transport = mock.create_autospec( big_query_read_grpc_transport.BigQueryReadGrpcTransport ) mock_client._create_bqstorage_client.return_value = bqstorage_client - session = storage.types.ReadSession() + session = bigquery_storage.types.ReadSession() bqstorage_client.create_read_session.return_value = session row_iterator = mut.RowIterator( mock_client, @@ -1981,7 +1987,7 @@ def test_to_arrow_w_bqstorage_creates_client(self): ) row_iterator.to_arrow(create_bqstorage_client=True) mock_client._create_bqstorage_client.assert_called_once() - bqstorage_client.transport.channel.close.assert_called_once() + bqstorage_client._transport.grpc_channel.close.assert_called_once() @unittest.skipIf(pyarrow is None, "Requires `pyarrow`") def test_to_arrow_create_bqstorage_client_wo_bqstorage(self): @@ -2011,13 +2017,15 @@ def test_to_arrow_create_bqstorage_client_wo_bqstorage(self): self.assertEqual(tbl.num_rows, 2) @unittest.skipIf(pyarrow is None, "Requires `pyarrow`") - @unittest.skipIf(storage is None, "Requires `google-cloud-bigquery-storage`") + @unittest.skipIf( + bigquery_storage is None, "Requires `google-cloud-bigquery-storage`" + ) def test_to_arrow_w_bqstorage_no_streams(self): from google.cloud.bigquery import schema from google.cloud.bigquery import table as mut - bqstorage_client = mock.create_autospec(storage.BigQueryReadClient) - session = storage.types.ReadSession() + bqstorage_client = mock.create_autospec(bigquery_storage.BigQueryReadClient) + session = bigquery_storage.types.ReadSession() arrow_schema = pyarrow.schema( [ pyarrow.field("colA", pyarrow.string()), @@ -2141,7 +2149,9 @@ def test_to_dataframe_iterable(self): self.assertEqual(df_2["age"][0], 33) @unittest.skipIf(pandas is None, "Requires `pandas`") - @unittest.skipIf(storage is None, "Requires `google-cloud-bigquery-storage`") + @unittest.skipIf( + bigquery_storage is None, "Requires `google-cloud-bigquery-storage`" + ) @unittest.skipIf(pyarrow is None, "Requires `pyarrow`") def test_to_dataframe_iterable_w_bqstorage(self): from google.cloud.bigquery import schema @@ -2156,8 +2166,8 @@ def test_to_dataframe_iterable_w_bqstorage(self): ] arrow_schema = pyarrow.schema(arrow_fields) - bqstorage_client = mock.create_autospec(storage.BigQueryReadClient) - bqstorage_client.transport = mock.create_autospec( + bqstorage_client = mock.create_autospec(bigquery_storage.BigQueryReadClient) + bqstorage_client._transport = mock.create_autospec( big_query_read_grpc_transport.BigQueryReadGrpcTransport ) streams = [ @@ -2165,7 +2175,7 @@ def test_to_dataframe_iterable_w_bqstorage(self): {"name": "/projects/proj/dataset/dset/tables/tbl/streams/1234"}, {"name": "/projects/proj/dataset/dset/tables/tbl/streams/5678"}, ] - session = storage.types.ReadSession( + session = bigquery_storage.types.ReadSession( streams=streams, arrow_schema={"serialized_schema": arrow_schema.serialize().to_pybytes()}, ) @@ -2208,7 +2218,7 @@ def test_to_dataframe_iterable_w_bqstorage(self): self.assertEqual(len(got), total_pages) # Don't close the client if it was passed in. - bqstorage_client.transport.channel.close.assert_not_called() + bqstorage_client._transport.grpc_channel.close.assert_not_called() @mock.patch("google.cloud.bigquery.table.pandas", new=None) def test_to_dataframe_iterable_error_if_pandas_is_none(self): @@ -2772,18 +2782,20 @@ def test_to_dataframe_max_results_w_create_bqstorage_warning(self): mock_client._create_bqstorage_client.assert_not_called() @unittest.skipIf(pandas is None, "Requires `pandas`") - @unittest.skipIf(storage is None, "Requires `google-cloud-bigquery-storage`") + @unittest.skipIf( + bigquery_storage is None, "Requires `google-cloud-bigquery-storage`" + ) def test_to_dataframe_w_bqstorage_creates_client(self): from google.cloud.bigquery import schema from google.cloud.bigquery import table as mut mock_client = _mock_client() - bqstorage_client = mock.create_autospec(storage.BigQueryReadClient) - bqstorage_client.transport = mock.create_autospec( + bqstorage_client = mock.create_autospec(bigquery_storage.BigQueryReadClient) + bqstorage_client._transport = mock.create_autospec( big_query_read_grpc_transport.BigQueryReadGrpcTransport ) mock_client._create_bqstorage_client.return_value = bqstorage_client - session = storage.types.ReadSession() + session = bigquery_storage.types.ReadSession() bqstorage_client.create_read_session.return_value = session row_iterator = mut.RowIterator( mock_client, @@ -2798,16 +2810,18 @@ def test_to_dataframe_w_bqstorage_creates_client(self): ) row_iterator.to_dataframe(create_bqstorage_client=True) mock_client._create_bqstorage_client.assert_called_once() - bqstorage_client.transport.channel.close.assert_called_once() + bqstorage_client._transport.grpc_channel.close.assert_called_once() @unittest.skipIf(pandas is None, "Requires `pandas`") - @unittest.skipIf(storage is None, "Requires `google-cloud-bigquery-storage`") + @unittest.skipIf( + bigquery_storage is None, "Requires `google-cloud-bigquery-storage`" + ) def test_to_dataframe_w_bqstorage_no_streams(self): from google.cloud.bigquery import schema from google.cloud.bigquery import table as mut - bqstorage_client = mock.create_autospec(storage.BigQueryReadClient) - session = storage.types.ReadSession() + bqstorage_client = mock.create_autospec(bigquery_storage.BigQueryReadClient) + session = bigquery_storage.types.ReadSession() bqstorage_client.create_read_session.return_value = session row_iterator = mut.RowIterator( @@ -2827,14 +2841,16 @@ def test_to_dataframe_w_bqstorage_no_streams(self): self.assertEqual(list(got), column_names) self.assertTrue(got.empty) - @unittest.skipIf(storage is None, "Requires `google-cloud-bigquery-storage`") + @unittest.skipIf( + bigquery_storage is None, "Requires `google-cloud-bigquery-storage`" + ) @unittest.skipIf(pandas is None, "Requires `pandas`") @unittest.skipIf(pyarrow is None, "Requires `pyarrow`") def test_to_dataframe_w_bqstorage_logs_session(self): from google.cloud.bigquery.table import Table - bqstorage_client = mock.create_autospec(storage.BigQueryReadClient) - session = storage.types.ReadSession() + bqstorage_client = mock.create_autospec(bigquery_storage.BigQueryReadClient) + session = bigquery_storage.types.ReadSession() session.name = "projects/test-proj/locations/us/sessions/SOMESESSION" bqstorage_client.create_read_session.return_value = session mock_logger = mock.create_autospec(logging.Logger) @@ -2851,7 +2867,9 @@ def test_to_dataframe_w_bqstorage_logs_session(self): ) @unittest.skipIf(pandas is None, "Requires `pandas`") - @unittest.skipIf(storage is None, "Requires `google-cloud-bigquery-storage`") + @unittest.skipIf( + bigquery_storage is None, "Requires `google-cloud-bigquery-storage`" + ) @unittest.skipIf(pyarrow is None, "Requires `pyarrow`") def test_to_dataframe_w_bqstorage_empty_streams(self): from google.cloud.bigquery import schema @@ -2866,8 +2884,8 @@ def test_to_dataframe_w_bqstorage_empty_streams(self): ] arrow_schema = pyarrow.schema(arrow_fields) - bqstorage_client = mock.create_autospec(storage.BigQueryReadClient) - session = storage.types.ReadSession( + bqstorage_client = mock.create_autospec(bigquery_storage.BigQueryReadClient) + session = bigquery_storage.types.ReadSession( streams=[{"name": "/projects/proj/dataset/dset/tables/tbl/streams/1234"}], arrow_schema={"serialized_schema": arrow_schema.serialize().to_pybytes()}, ) @@ -2904,7 +2922,9 @@ def test_to_dataframe_w_bqstorage_empty_streams(self): self.assertTrue(got.empty) @unittest.skipIf(pandas is None, "Requires `pandas`") - @unittest.skipIf(storage is None, "Requires `google-cloud-bigquery-storage`") + @unittest.skipIf( + bigquery_storage is None, "Requires `google-cloud-bigquery-storage`" + ) @unittest.skipIf(pyarrow is None, "Requires `pyarrow`") def test_to_dataframe_w_bqstorage_nonempty(self): from google.cloud.bigquery import schema @@ -2919,8 +2939,8 @@ def test_to_dataframe_w_bqstorage_nonempty(self): ] arrow_schema = pyarrow.schema(arrow_fields) - bqstorage_client = mock.create_autospec(storage.BigQueryReadClient) - bqstorage_client.transport = mock.create_autospec( + bqstorage_client = mock.create_autospec(bigquery_storage.BigQueryReadClient) + bqstorage_client._transport = mock.create_autospec( big_query_read_grpc_transport.BigQueryReadGrpcTransport ) streams = [ @@ -2928,7 +2948,7 @@ def test_to_dataframe_w_bqstorage_nonempty(self): {"name": "/projects/proj/dataset/dset/tables/tbl/streams/1234"}, {"name": "/projects/proj/dataset/dset/tables/tbl/streams/5678"}, ] - session = storage.types.ReadSession( + session = bigquery_storage.types.ReadSession( streams=streams, arrow_schema={"serialized_schema": arrow_schema.serialize().to_pybytes()}, ) @@ -2979,10 +2999,12 @@ def test_to_dataframe_w_bqstorage_nonempty(self): self.assertEqual(len(got.index), total_rows) # Don't close the client if it was passed in. - bqstorage_client.transport.channel.close.assert_not_called() + bqstorage_client._transport.grpc_channel.close.assert_not_called() @unittest.skipIf(pandas is None, "Requires `pandas`") - @unittest.skipIf(storage is None, "Requires `google-cloud-bigquery-storage`") + @unittest.skipIf( + bigquery_storage is None, "Requires `google-cloud-bigquery-storage`" + ) @unittest.skipIf(pyarrow is None, "Requires `pyarrow`") def test_to_dataframe_w_bqstorage_multiple_streams_return_unique_index(self): from google.cloud.bigquery import schema @@ -2996,12 +3018,12 @@ def test_to_dataframe_w_bqstorage_multiple_streams_return_unique_index(self): {"name": "/projects/proj/dataset/dset/tables/tbl/streams/1234"}, {"name": "/projects/proj/dataset/dset/tables/tbl/streams/5678"}, ] - session = storage.types.ReadSession( + session = bigquery_storage.types.ReadSession( streams=streams, arrow_schema={"serialized_schema": arrow_schema.serialize().to_pybytes()}, ) - bqstorage_client = mock.create_autospec(storage.BigQueryReadClient) + bqstorage_client = mock.create_autospec(bigquery_storage.BigQueryReadClient) bqstorage_client.create_read_session.return_value = session mock_rowstream = mock.create_autospec(reader.ReadRowsStream) @@ -3034,7 +3056,9 @@ def test_to_dataframe_w_bqstorage_multiple_streams_return_unique_index(self): self.assertTrue(got.index.is_unique) @unittest.skipIf(pandas is None, "Requires `pandas`") - @unittest.skipIf(storage is None, "Requires `google-cloud-bigquery-storage`") + @unittest.skipIf( + bigquery_storage is None, "Requires `google-cloud-bigquery-storage`" + ) @unittest.skipIf(tqdm is None, "Requires `tqdm`") @unittest.skipIf(pyarrow is None, "Requires `pyarrow`") @mock.patch("tqdm.tqdm") @@ -3049,14 +3073,14 @@ def test_to_dataframe_w_bqstorage_updates_progress_bar(self, tqdm_mock): arrow_fields = [pyarrow.field("testcol", pyarrow.int64())] arrow_schema = pyarrow.schema(arrow_fields) - bqstorage_client = mock.create_autospec(storage.BigQueryReadClient) + bqstorage_client = mock.create_autospec(bigquery_storage.BigQueryReadClient) streams = [ # Use two streams we want to check that progress bar updates are # sent from each stream. {"name": "/projects/proj/dataset/dset/tables/tbl/streams/1234"}, {"name": "/projects/proj/dataset/dset/tables/tbl/streams/5678"}, ] - session = storage.types.ReadSession( + session = bigquery_storage.types.ReadSession( streams=streams, arrow_schema={"serialized_schema": arrow_schema.serialize().to_pybytes()}, ) @@ -3111,7 +3135,9 @@ def blocking_to_arrow(*args, **kwargs): tqdm_mock().close.assert_called_once() @unittest.skipIf(pandas is None, "Requires `pandas`") - @unittest.skipIf(storage is None, "Requires `google-cloud-bigquery-storage`") + @unittest.skipIf( + bigquery_storage is None, "Requires `google-cloud-bigquery-storage`" + ) @unittest.skipIf(pyarrow is None, "Requires `pyarrow`") def test_to_dataframe_w_bqstorage_exits_on_keyboardinterrupt(self): from google.cloud.bigquery import schema @@ -3129,8 +3155,8 @@ def test_to_dataframe_w_bqstorage_exits_on_keyboardinterrupt(self): ] arrow_schema = pyarrow.schema(arrow_fields) - bqstorage_client = mock.create_autospec(storage.BigQueryReadClient) - session = storage.types.ReadSession( + bqstorage_client = mock.create_autospec(bigquery_storage.BigQueryReadClient) + session = bigquery_storage.types.ReadSession( streams=[ # Use multiple streams because one will fail with a # KeyboardInterrupt, and we want to check that the other streams @@ -3228,11 +3254,13 @@ def test_to_dataframe_tabledata_list_w_multiple_pages_return_unique_index(self): self.assertTrue(df.index.is_unique) @unittest.skipIf(pandas is None, "Requires `pandas`") - @unittest.skipIf(storage is None, "Requires `google-cloud-bigquery-storage`") + @unittest.skipIf( + bigquery_storage is None, "Requires `google-cloud-bigquery-storage`" + ) def test_to_dataframe_w_bqstorage_raises_auth_error(self): from google.cloud.bigquery import table as mut - bqstorage_client = mock.create_autospec(storage.BigQueryReadClient) + bqstorage_client = mock.create_autospec(bigquery_storage.BigQueryReadClient) bqstorage_client.create_read_session.side_effect = google.api_core.exceptions.Forbidden( "TEST BigQuery Storage API not enabled. TEST" ) @@ -3245,12 +3273,14 @@ def test_to_dataframe_w_bqstorage_raises_auth_error(self): with pytest.raises(google.api_core.exceptions.Forbidden): row_iterator.to_dataframe(bqstorage_client=bqstorage_client) - @unittest.skipIf(storage is None, "Requires `google-cloud-bigquery-storage`") + @unittest.skipIf( + bigquery_storage is None, "Requires `google-cloud-bigquery-storage`" + ) def test_to_dataframe_w_bqstorage_partition(self): from google.cloud.bigquery import schema from google.cloud.bigquery import table as mut - bqstorage_client = mock.create_autospec(storage.BigQueryReadClient) + bqstorage_client = mock.create_autospec(bigquery_storage.BigQueryReadClient) row_iterator = mut.RowIterator( _mock_client(), @@ -3263,12 +3293,14 @@ def test_to_dataframe_w_bqstorage_partition(self): with pytest.raises(ValueError): row_iterator.to_dataframe(bqstorage_client) - @unittest.skipIf(storage is None, "Requires `google-cloud-bigquery-storage`") + @unittest.skipIf( + bigquery_storage is None, "Requires `google-cloud-bigquery-storage`" + ) def test_to_dataframe_w_bqstorage_snapshot(self): from google.cloud.bigquery import schema from google.cloud.bigquery import table as mut - bqstorage_client = mock.create_autospec(storage.BigQueryReadClient) + bqstorage_client = mock.create_autospec(bigquery_storage.BigQueryReadClient) row_iterator = mut.RowIterator( _mock_client(), @@ -3282,7 +3314,9 @@ def test_to_dataframe_w_bqstorage_snapshot(self): row_iterator.to_dataframe(bqstorage_client) @unittest.skipIf(pandas is None, "Requires `pandas`") - @unittest.skipIf(storage is None, "Requires `google-cloud-bigquery-storage`") + @unittest.skipIf( + bigquery_storage is None, "Requires `google-cloud-bigquery-storage`" + ) @unittest.skipIf(pyarrow is None, "Requires `pyarrow`") def test_to_dataframe_concat_categorical_dtype_w_pyarrow(self): from google.cloud.bigquery import schema @@ -3300,11 +3334,11 @@ def test_to_dataframe_concat_categorical_dtype_w_pyarrow(self): arrow_schema = pyarrow.schema(arrow_fields) # create a mock BQ storage client - bqstorage_client = mock.create_autospec(storage.BigQueryReadClient) - bqstorage_client.transport = mock.create_autospec( + bqstorage_client = mock.create_autospec(bigquery_storage.BigQueryReadClient) + bqstorage_client._transport = mock.create_autospec( big_query_read_grpc_transport.BigQueryReadGrpcTransport ) - session = storage.types.ReadSession( + session = bigquery_storage.types.ReadSession( streams=[{"name": "/projects/proj/dataset/dset/tables/tbl/streams/1234"}], arrow_schema={"serialized_schema": arrow_schema.serialize().to_pybytes()}, ) @@ -3388,7 +3422,7 @@ def test_to_dataframe_concat_categorical_dtype_w_pyarrow(self): ) # Don't close the client if it was passed in. - bqstorage_client.transport.channel.close.assert_not_called() + bqstorage_client._transport.grpc_channel.close.assert_not_called() @unittest.skipIf(pandas is None, "Requires `pandas`") def test_to_dataframe_concat_categorical_dtype_wo_pyarrow(self): @@ -3830,7 +3864,9 @@ def test_set_expiration_w_none(self): assert time_partitioning._properties["expirationMs"] is None -@pytest.mark.skipif(storage is None, reason="Requires `google-cloud-bigquery-storage`") +@pytest.mark.skipif( + bigquery_storage is None, reason="Requires `google-cloud-bigquery-storage`" +) @pytest.mark.parametrize( "table_path", (