diff --git a/packages/google-cloud-bigquery/google/cloud/bigquery/_arrow_query.py b/packages/google-cloud-bigquery/google/cloud/bigquery/_arrow_query.py index f627e4eb619b..30534cbe4a3e 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery/_arrow_query.py +++ b/packages/google-cloud-bigquery/google/cloud/bigquery/_arrow_query.py @@ -160,23 +160,13 @@ def download_arrow_from_job( if pyarrow is None: raise ValueError(_NO_PYARROW_ERROR) - # Step 1: Ensure BigQuery Read API client is available upfront. - if bqstorage_client is None: - if row_iterator.client is None: - raise ValueError("RowIterator client is None.") - bqstorage_client = row_iterator.client._ensure_bqstorage_client() - if bqstorage_client is None: - raise ValueError( - "The google-cloud-bigquery-storage library is required to read Arrow results." - ) - offset = 0 pa_schema = None total_rows = row_iterator.total_rows job_complete = False remaining = row_iterator.max_results - # Step 2: Process initial inline Arrow response from jobs.query if available. + # Step 1: Process initial inline Arrow response from jobs.query if available. if row_iterator._first_page_response: first_page = row_iterator._first_page_response @@ -206,12 +196,22 @@ def download_arrow_from_job( batch, remaining = _slice_batch(batch, remaining) yield batch - # Step 3: Return early if all results were delivered in the first page + # Step 2: Return early if all results were delivered in the first page # response, or if the first page already reached max_results. if (job_complete and offset >= total_rows) or remaining == 0: return - # Step 4: Stream remaining Arrow record batches from the job default stream. + # Step 3: Stream remaining Arrow record batches from the job default stream, + # with the BigQuery client's shared BigQuery Storage client unless the + # caller passed one. Results that fit in the first page never create it. + if bqstorage_client is None: + if row_iterator.client is None: + raise ValueError("RowIterator client is None.") + bqstorage_client = row_iterator.client._get_shared_bqstorage_client() + if bqstorage_client is None: + raise ValueError( + "The google-cloud-bigquery-storage library is required to read Arrow results." + ) project = row_iterator._project or ( row_iterator.client.project if row_iterator.client else None ) diff --git a/packages/google-cloud-bigquery/google/cloud/bigquery/client.py b/packages/google-cloud-bigquery/google/cloud/bigquery/client.py index b78438c5bc23..46f9b638ff2a 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery/client.py +++ b/packages/google-cloud-bigquery/google/cloud/bigquery/client.py @@ -26,6 +26,7 @@ import math import os import tempfile +import threading import typing import uuid import warnings @@ -166,6 +167,35 @@ # https://github.com/googleapis/python-bigquery/issues/438 _MIN_GET_QUERY_RESULTS_TIMEOUT = 120 +# While a stream is open and the connection is idle, the shared BigQuery +# Storage client pings the server after a minute, and closes the connection +# if a ping isn't answered within a minute. A dropped ReadRows stream is +# resumed at the next unread row. +_BQSTORAGE_KEEPALIVE_TIME_MS = 60 * 1000 +_BQSTORAGE_KEEPALIVE_TIMEOUT_MS = 60 * 1000 + + +def _create_bqstorage_transport(**kwargs): + """Create a BigQuery Storage gRPC transport whose channel sends keepalive + pings. + + ``BigQueryReadClient`` calls this with the arguments it passes to its + default transport. + """ + from google.cloud import bigquery_storage + + transport_class = bigquery_storage.BigQueryReadClient.get_transport_class("grpc") + + def create_channel(*args, options=(), **channel_kwargs): + options = list(options) + [ + ("grpc.keepalive_time_ms", _BQSTORAGE_KEEPALIVE_TIME_MS), + ("grpc.keepalive_timeout_ms", _BQSTORAGE_KEEPALIVE_TIMEOUT_MS), + ] + return transport_class.create_channel(*args, options=options, **channel_kwargs) + + return transport_class(channel=create_channel, **kwargs) + + _LOAD_TABLE_FROM_DATAFRAME_DEPRECATED = ( "Loading DataFrames via google-cloud-bigquery is deprecated. " "For direct, optimized loading, please call 'pandas_gbq.to_gbq()' directly." @@ -300,6 +330,14 @@ def __init__( self._default_load_job_config = copy.deepcopy(default_load_job_config) self.default_job_creation_mode = default_job_creation_mode + # ARROW query results share one BigQuery Storage client. It's created + # the first time a query needs a stream, and close() closes it. + self._universe_domain = _get_client_universe(client_options) + self._shared_bqstorage_client: Optional[ + "google.cloud.bigquery_storage.BigQueryReadClient" + ] = None + self._shared_bqstorage_client_lock = threading.Lock() + # Use property setter so validation can run. self.default_query_job_config = default_query_job_config @@ -348,6 +386,9 @@ def default_load_job_config(self, value: LoadJobConfig): def close(self): """Close the underlying transport objects, releasing system resources. + This also closes the BigQuery Storage client that ``ARROW`` query + results share. + .. note:: The client instance can be used for making additional requests even @@ -356,6 +397,11 @@ def close(self): """ self._http._auth_request.session.close() self._http.close() + with self._shared_bqstorage_client_lock: + bqstorage_client = self._shared_bqstorage_client + self._shared_bqstorage_client = None + if bqstorage_client is not None: + bqstorage_client._transport.close() def get_service_account_email( self, @@ -600,6 +646,7 @@ def _ensure_bqstorage_client( client_info: Optional[ "google.api_core.gapic_v1.client_info.ClientInfo" ] = DEFAULT_BQSTORAGE_CLIENT_INFO, + transport: Optional[Callable[..., Any]] = None, ) -> Optional["google.cloud.bigquery_storage.BigQueryReadClient"]: """Create a BigQuery Storage API client using this client's credentials. @@ -613,6 +660,9 @@ def _ensure_bqstorage_client( client_info: The client info used with a new BigQuery Storage client instance if one is created. + transport: + A callable that creates the transport of a new BigQuery Storage + client instance. If ``None``, the default transport is used. Returns: A BigQuery Storage API client. @@ -680,10 +730,37 @@ def _ensure_bqstorage_client( credentials=self._credentials, client_options=client_options, client_info=amended_client_info, + transport=transport, ) return bqstorage_client + def _get_shared_bqstorage_client( + self, + ) -> Optional["google.cloud.bigquery_storage.BigQueryReadClient"]: + """Return the BigQuery Storage client that ``ARROW`` query results share. + + The client is created on first use, with this client's credentials and + universe domain and a gRPC channel that sends keepalive pings. + :meth:`close` closes it, and the next call creates a new one. + + Returns: + A BigQuery Storage API client, or ``None`` if + google-cloud-bigquery-storage is missing or outdated. + """ + with self._shared_bqstorage_client_lock: + if self._shared_bqstorage_client is None: + client_options = None + if self._universe_domain != _DEFAULT_UNIVERSE: + client_options = google.api_core.client_options.ClientOptions( + universe_domain=self._universe_domain + ) + self._shared_bqstorage_client = self._ensure_bqstorage_client( + client_options=client_options, + transport=_create_bqstorage_transport, + ) + return self._shared_bqstorage_client + def _dataset_from_arg(self, dataset) -> Union[Dataset, DatasetReference]: if isinstance(dataset, str): dataset = DatasetReference.from_string( diff --git a/packages/google-cloud-bigquery/google/cloud/bigquery/table.py b/packages/google-cloud-bigquery/google/cloud/bigquery/table.py index f730914b16b9..b7e2bd55fcea 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery/table.py +++ b/packages/google-cloud-bigquery/google/cloud/bigquery/table.py @@ -2317,7 +2317,8 @@ def _download_arrow_from_job_id( Args: bqstorage_client (Optional[bigquery_storage.BigQueryReadClient]): - Client for BigQuery Storage Read API. If None, one will be created. + Client for BigQuery Storage Read API. If None, the BigQuery + Storage client that the BigQuery client shares is used. timeout (Optional[float]): Timeout in seconds for read operations. diff --git a/packages/google-cloud-bigquery/tests/system/test_arrow.py b/packages/google-cloud-bigquery/tests/system/test_arrow.py index cc3ed5710b9a..339e6475a1ac 100644 --- a/packages/google-cloud-bigquery/tests/system/test_arrow.py +++ b/packages/google-cloud-bigquery/tests/system/test_arrow.py @@ -755,3 +755,69 @@ def test_query_and_wait_arrow_max_results( table = pyarrow.Table.from_batches(batches) assert table.column("x").to_pylist() == list(range(1, expected_rows + 1)) assert read_rows.called == streamed + + +def _query_arrow(client, row_count): + return client.query_and_wait( + f"SELECT x FROM UNNEST(GENERATE_ARRAY(1, {row_count})) AS x ORDER BY x", + job_config=bigquery.QueryJobConfig(use_query_cache=False), + page_size=100, + query_results_format=enums.QueryResultsFormat.ARROW, + ) + + +def test_query_and_wait_arrow_queries_share_one_bqstorage_client(): + from google.cloud.bigquery import client as client_module + + with bigquery.Client() as client: + # Wrap the client factory to count how many clients get created. + ensure = mock.Mock(wraps=client._ensure_bqstorage_client) + with mock.patch.object(client, "_ensure_bqstorage_client", ensure): + tables = [ + pyarrow.Table.from_batches( + list(_query_arrow(client, 1000).to_arrow_iterable()) + ), + _query_arrow(client, 1000).to_arrow(), + pyarrow.Table.from_batches( + list(_query_arrow(client, 1000).to_arrow_iterable()) + ), + ] + shared = client._shared_bqstorage_client + + for table in tables: + assert table.column("x").to_pylist() == list(range(1, 1001)) + ensure.assert_called_once_with( + client_options=None, transport=client_module._create_bqstorage_transport + ) + assert shared is not None + # Leaving the with block calls close(), which closes the shared channel. + assert client._shared_bqstorage_client is None + with pytest.raises(ValueError, match="(?i)closed"): + shared.read_rows("projects/p/locations/US/jobs/j/streams/_default") + + +def test_query_and_wait_arrow_ending_a_read_early_cancels_only_its_stream(): + from google.cloud.bigquery_storage_v1 import reader + + # Wrap ReadRowsStream.cancel() to count calls without changing it. + cancel_stream = reader.ReadRowsStream.cancel + with bigquery.Client() as client, mock.patch.object( + reader.ReadRowsStream, "cancel", autospec=True, side_effect=cancel_stream + ) as cancel: + ended_early_rows = _query_arrow(client, 100000) + ended_early = ended_early_rows.to_arrow_iterable() + read_fully = _query_arrow(client, 100000).to_arrow_iterable() + # Two reads reach the stream, whether or not jobs.query returned a + # first page. + next(ended_early) + next(ended_early) + rows = [next(read_fully), next(read_fully)] + + ended_early.close() + rows.extend(read_fully) + + (stream,) = cancel.call_args.args + cancel.assert_called_once_with(stream) + assert stream._name.endswith(f"/jobs/{ended_early_rows.job_id}/streams/_default") + table = pyarrow.Table.from_batches(rows) + assert table.column("x").to_pylist() == list(range(1, 100001)) diff --git a/packages/google-cloud-bigquery/tests/unit/test_client.py b/packages/google-cloud-bigquery/tests/unit/test_client.py index 88629500a007..7f49ddfadc33 100644 --- a/packages/google-cloud-bigquery/tests/unit/test_client.py +++ b/packages/google-cloud-bigquery/tests/unit/test_client.py @@ -1033,6 +1033,128 @@ def test_ensure_bqstorage_client_existing_client_check_fails(self): ] assert matching_warnings, "Obsolete dependency warning not raised." + def test_get_shared_bqstorage_client_creates_one_client(self): + bigquery_storage = pytest.importorskip("google.cloud.bigquery_storage") + from google.cloud.bigquery import client as client_module + + mock_client = mock.create_autospec(bigquery_storage.BigQueryReadClient) + creds = _make_credentials() + client = self._make_one(project=self.PROJECT, credentials=creds) + + with mock.patch( + "google.cloud.bigquery_storage.BigQueryReadClient", mock_client + ): + first = client._get_shared_bqstorage_client() + second = client._get_shared_bqstorage_client() + + self.assertIs(first, mock_client.return_value) + self.assertIs(second, first) + mock_client.assert_called_once_with( + credentials=creds, + client_options=None, + client_info=mock.ANY, + transport=client_module._create_bqstorage_transport, + ) + + @pytest.mark.skipif( + packaging.version.parse(getattr(google.api_core, "__version__", "0.0.0")) + < packaging.version.Version("2.15.0"), + reason="universe_domain not supported with google-api-core < 2.15.0", + ) + def test_get_shared_bqstorage_client_uses_universe_domain(self): + bigquery_storage = pytest.importorskip("google.cloud.bigquery_storage") + + mock_client = mock.create_autospec(bigquery_storage.BigQueryReadClient) + creds = _make_credentials() + client = self._make_one( + project=self.PROJECT, + credentials=creds, + client_options={"universe_domain": "foo.com"}, + ) + + with mock.patch( + "google.cloud.bigquery_storage.BigQueryReadClient", mock_client + ): + client._get_shared_bqstorage_client() + + _, kwargs = mock_client.call_args + self.assertEqual(kwargs["client_options"].universe_domain, "foo.com") + + def test_get_shared_bqstorage_client_tries_again_if_unavailable(self): + from google.cloud.bigquery import client as client_module + + creds = _make_credentials() + client = self._make_one(project=self.PROJECT, credentials=creds) + + with mock.patch.object( + client, "_ensure_bqstorage_client", return_value=None + ) as ensure: + self.assertIsNone(client._get_shared_bqstorage_client()) + self.assertIsNone(client._get_shared_bqstorage_client()) + + expected = mock.call( + client_options=None, transport=client_module._create_bqstorage_transport + ) + self.assertEqual(ensure.call_args_list, [expected, expected]) + + def test_get_shared_bqstorage_client_creates_one_client_across_threads(self): + import threading + import time + + from google.cloud.bigquery import client as client_module + + creds = _make_credentials() + client = self._make_one(project=self.PROJECT, credentials=creds) + + def slow_create(**kwargs): + time.sleep(0.05) + return mock.Mock() + + results = [] + + def get_client(): + results.append(client._get_shared_bqstorage_client()) + + with mock.patch.object( + client, "_ensure_bqstorage_client", side_effect=slow_create + ) as ensure: + threads = [threading.Thread(target=get_client) for _ in range(5)] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + + ensure.assert_called_once_with( + client_options=None, transport=client_module._create_bqstorage_transport + ) + self.assertEqual(len(results), 5) + self.assertTrue(all(result is results[0] for result in results)) + + def test_get_shared_bqstorage_client_channel_sends_keepalive_pings(self): + bigquery_storage = pytest.importorskip("google.cloud.bigquery_storage") + import google.auth.credentials + import grpc + + transport_class = bigquery_storage.BigQueryReadClient.get_transport_class( + "grpc" + ) + creds = google.auth.credentials.AnonymousCredentials() + client = self._make_one(project=self.PROJECT, credentials=creds) + channel = grpc.insecure_channel("localhost:1") + + with mock.patch.object( + transport_class, "create_channel", return_value=channel + ) as create_channel: + bqstorage_client = client._get_shared_bqstorage_client() + channel.close() + + self.assertIsInstance(bqstorage_client._transport, transport_class) + options = create_channel.call_args.kwargs["options"] + self.assertIn(("grpc.keepalive_time_ms", 60000), options) + self.assertIn(("grpc.keepalive_timeout_ms", 60000), options) + # The transport's own options are kept. + self.assertIn(("grpc.max_receive_message_length", -1), options) + def test_create_routine_w_minimal_resource(self): from google.cloud.bigquery.routine import Routine, RoutineReference @@ -1653,6 +1775,23 @@ def test_close(self): http.close.assert_called_once() http._auth_request.session.close.assert_called_once() + def test_close_closes_shared_bqstorage_client(self): + creds = _make_credentials() + http = mock.Mock() + http._auth_request.session = mock.Mock() + client = self._make_one(project=self.PROJECT, credentials=creds, _http=http) + first, second = mock.Mock(), mock.Mock() + + with mock.patch.object( + client, "_ensure_bqstorage_client", side_effect=[first, second] + ): + self.assertIs(client._get_shared_bqstorage_client(), first) + client.close() + self.assertIs(client._get_shared_bqstorage_client(), second) + + first._transport.close.assert_called_once_with() + second._transport.close.assert_not_called() + def test_get_model(self): path = "projects/%s/datasets/%s/models/%s" % ( self.PROJECT, diff --git a/packages/google-cloud-bigquery/tests/unit/test_query_results_format_arrow.py b/packages/google-cloud-bigquery/tests/unit/test_query_results_format_arrow.py index 81d3e6abbf30..46489b355585 100644 --- a/packages/google-cloud-bigquery/tests/unit/test_query_results_format_arrow.py +++ b/packages/google-cloud-bigquery/tests/unit/test_query_results_format_arrow.py @@ -211,7 +211,7 @@ def test_row_iterator_to_arrow_iterable_delegates_when_format_is_arrow(self): def test_download_arrow_from_job_id_constructs_stream_and_reads(self): mock_client = mock.MagicMock() mock_bqstorage = mock.MagicMock() - mock_client._ensure_bqstorage_client.return_value = mock_bqstorage + mock_client._get_shared_bqstorage_client.return_value = mock_bqstorage iterator = RowIterator( client=mock_client, @@ -247,7 +247,7 @@ def test_download_arrow_from_job_id_constructs_stream_and_reads(self): def test_download_arrow_from_job_id_with_first_page_response(self): mock_client = mock.MagicMock() mock_bqstorage = mock.MagicMock() - mock_client._ensure_bqstorage_client.return_value = mock_bqstorage + mock_client._get_shared_bqstorage_client.return_value = mock_bqstorage raw_schema_bytes = b"schema_bytes_123" raw_batch_bytes = b"batch_bytes_123" @@ -321,7 +321,7 @@ def test_download_arrow_from_job_id_with_first_page_response(self): def test_download_arrow_from_job_id_with_first_page_response_schema_only(self): mock_client = mock.MagicMock() mock_bqstorage = mock.MagicMock() - mock_client._ensure_bqstorage_client.return_value = mock_bqstorage + mock_client._get_shared_bqstorage_client.return_value = mock_bqstorage raw_schema_bytes = b"schema_bytes_456" first_page_response = { @@ -373,7 +373,7 @@ def test_download_arrow_from_job_id_with_first_page_response_schema_only(self): def test_download_arrow_from_job_id_missing_storage_client_raises_value_error(self): mock_client = mock.MagicMock() - mock_client._ensure_bqstorage_client.return_value = None + mock_client._get_shared_bqstorage_client.return_value = None iterator = RowIterator( client=mock_client, @@ -397,7 +397,7 @@ def test_download_arrow_from_job_id_missing_storage_client_raises_value_error(se def test_download_arrow_from_job_id_with_schema_and_batch(self): mock_client = mock.MagicMock() mock_bqstorage = mock.MagicMock() - mock_client._ensure_bqstorage_client.return_value = mock_bqstorage + mock_client._get_shared_bqstorage_client.return_value = mock_bqstorage iterator = RowIterator( client=mock_client, @@ -476,14 +476,14 @@ def test_download_arrow_from_job_id_avoids_read_rows_when_all_rows_present(self) batches = list(iterator._download_arrow_from_job_id(timeout=5.0)) self.assertEqual(batches, [mock_first_batch]) - mock_client._ensure_bqstorage_client.assert_called_once() + mock_client._get_shared_bqstorage_client.assert_not_called() batches_second = list(iterator._download_arrow_from_job_id(timeout=5.0)) self.assertEqual(batches_second, [mock_first_batch]) def test_download_arrow_from_job_id_calls_read_rows_when_job_not_complete(self): mock_client = mock.MagicMock() mock_bqstorage = mock.MagicMock() - mock_client._ensure_bqstorage_client.return_value = mock_bqstorage + mock_client._get_shared_bqstorage_client.return_value = mock_bqstorage raw_schema_bytes = b"schema_bytes_789" raw_batch_bytes = b"batch_bytes_789" @@ -520,7 +520,7 @@ def test_download_arrow_from_job_id_calls_read_rows_when_job_not_complete(self): batches = list(iterator._download_arrow_from_job_id(timeout=5.0)) self.assertEqual(batches, [mock_first_batch]) - mock_client._ensure_bqstorage_client.assert_called_once() + mock_client._get_shared_bqstorage_client.assert_called_once_with() expected_stream = "projects/test-proj/locations/US/jobs/test-job-incomplete/streams/_default" mock_bqstorage.read_rows.assert_called_once_with( expected_stream, @@ -576,7 +576,7 @@ def test_job_helpers_query_and_wait_accepts_enums(self): def test_row_iterator_to_dataframe_iterable_when_format_is_arrow(self): pytest.importorskip("pandas") mock_client = mock.MagicMock() - mock_client._ensure_bqstorage_client.return_value = mock.MagicMock() + mock_client._get_shared_bqstorage_client.return_value = mock.MagicMock() iterator = RowIterator( client=mock_client, @@ -599,7 +599,7 @@ def test_row_iterator_to_dataframe_iterable_when_format_is_arrow(self): def test_row_iterator_to_dataframe_when_format_is_arrow(self): pytest.importorskip("pandas") mock_client = mock.MagicMock() - mock_client._ensure_bqstorage_client.return_value = mock.MagicMock() + mock_client._get_shared_bqstorage_client.return_value = mock.MagicMock() iterator = RowIterator( client=mock_client, @@ -1198,5 +1198,49 @@ def test_to_arrow_with_arrow_format_and_max_results_keeps_bqstorage_client(): assert not [w for w in caught if "max_results" in str(w.message)] +def test_download_arrow_uses_passed_client_and_leaves_it_open(): + iterator, bqstorage_client, _ = _limited_iterator([10, 15], 25, max_results=None) + + batches = list(iterator._download_arrow_from_job_id(bqstorage_client)) + + assert _values(batches) == list(range(1, 26)) + iterator.client._get_shared_bqstorage_client.assert_not_called() + bqstorage_client._transport.close.assert_not_called() + + +def test_download_arrow_needs_a_client_only_to_stream(): + iterator, _, _ = _limited_iterator([10, 15], 25, max_results=None) + iterator.client = None + batches = iterator._download_arrow_from_job_id() + + assert next(batches).num_rows == 10 + with pytest.raises(ValueError, match="RowIterator client is None"): + next(batches) + + +def test_to_arrow_with_arrow_format_reads_with_the_shared_client(): + pyarrow = pytest.importorskip("pyarrow") + client = mock.MagicMock() + iterator = RowIterator( + client=client, + api_request=mock.MagicMock(), + path=None, + schema=(), + query_results_format="ARROW", + ) + batch = pyarrow.record_batch([pyarrow.array([1, 2])], names=["x"]) + + with mock.patch.object( + iterator, "to_arrow_iterable", return_value=iter([batch]) + ) as to_arrow_iterable: + table = iterator.to_arrow() + + assert table.num_rows == 2 + # to_arrow() doesn't create (and then close) a client of its own; the + # stream uses the BigQuery client's shared one. + client._ensure_bqstorage_client.assert_not_called() + assert to_arrow_iterable.call_args.kwargs["bqstorage_client"] is None + + if __name__ == "__main__": unittest.main()