diff --git a/packages/google-cloud-bigquery-storage/google/cloud/bigquery_storage_v1/client.py b/packages/google-cloud-bigquery-storage/google/cloud/bigquery_storage_v1/client.py index 0decb4f413ba..791be3f2a0dc 100644 --- a/packages/google-cloud-bigquery-storage/google/cloud/bigquery_storage_v1/client.py +++ b/packages/google-cloud-bigquery-storage/google/cloud/bigquery_storage_v1/client.py @@ -57,6 +57,8 @@ def read_rows( timeout=google.api_core.gapic_v1.method.DEFAULT, metadata=(), retry_delay_callback=None, + *, + arrow_serialization_options=None, ): """ Reads rows from the table in the format prescribed by the read @@ -124,6 +126,14 @@ def read_rows( BigQueryReadClient will call retry_delay_callback with the delay duration (in seconds) before it starts sleeping until the next attempt. + arrow_serialization_options (Optional[Union[dict, \ + ~google.cloud.bigquery_storage_v1.types.ArrowSerializationOptions \ + ]]): + Arrow serialization options, such as buffer compression, to + send on every ``ReadRows`` request, including after a + reconnect. Only job default streams + (``projects/{project}/locations/{location}/jobs/{job}/streams/_default``) + accept them. Returns: ~google.cloud.bigquery_storage_v1.reader.ReadRowsStream: @@ -137,6 +147,11 @@ def read_rows( to a retryable error and retry attempts failed. ValueError: If the parameters are invalid. """ + request_template = None + if arrow_serialization_options is not None: + request_template = { + "arrow_serialization_options": arrow_serialization_options + } gapic_client = super(BigQueryReadClient, self) stream = reader.ReadRowsStream( gapic_client, @@ -144,6 +159,7 @@ def read_rows( offset, {"retry": retry, "timeout": timeout, "metadata": metadata}, retry_delay_callback=retry_delay_callback, + request_template=request_template, ) stream._reconnect() return stream diff --git a/packages/google-cloud-bigquery-storage/google/cloud/bigquery_storage_v1/reader.py b/packages/google-cloud-bigquery-storage/google/cloud/bigquery_storage_v1/reader.py index 27ab7f55de48..ade01abb67b8 100644 --- a/packages/google-cloud-bigquery-storage/google/cloud/bigquery_storage_v1/reader.py +++ b/packages/google-cloud-bigquery-storage/google/cloud/bigquery_storage_v1/reader.py @@ -27,6 +27,8 @@ import google.api_core.exceptions import google.rpc.error_details_pb2 # type: ignore +from google.cloud.bigquery_storage_v1 import types as gapic_types + try: import pandas except ImportError: # pragma: NO COVER @@ -86,7 +88,14 @@ class ReadRowsStream(object): """ def __init__( - self, client, name, offset, read_rows_kwargs, retry_delay_callback=None + self, + client, + name, + offset, + read_rows_kwargs, + retry_delay_callback=None, + *, + request_template=None, ): """Construct a ReadRowsStream. @@ -114,6 +123,13 @@ def __init__( ReadRowsStream will call retry_delay_callback with the delay duration (in seconds) before it starts sleeping until the next attempt. + request_template (Optional[Union[dict, \ + ~google.cloud.bigquery_storage_v1.types.ReadRowsRequest \ + ]]): + Request fields, such as ``arrow_serialization_options``, to + send on every ``ReadRows`` request, including after a + reconnect. ``read_stream`` and ``offset`` are set from + ``name`` and the current stream offset. Returns: Iterable[ \ @@ -129,6 +145,8 @@ def __init__( self._offset = offset self._read_rows_kwargs = read_rows_kwargs self._retry_delay_callback = retry_delay_callback + self._request_template = request_template + self._cancelled = False self._wrapped = None def __iter__(self): @@ -140,19 +158,20 @@ def __iter__(self): ]: A sequence of row messages. """ - # Infinite loop to reconnect on reconnectable errors while processing - # the row stream. + # Loop to reconnect on reconnectable errors while processing the row + # stream, until the stream ends or is cancelled. if self._wrapped is None: self._reconnect() - while True: + while not self._cancelled: try: for message in self._wrapped: rowcount = message.row_count self._offset += rowcount yield message + self._wrapped = None return # Made it through the whole stream. except google.api_core.exceptions.InternalServerError as exc: resumable_error = any( @@ -165,22 +184,60 @@ def __iter__(self): # Transient error, so reconnect to the stream. pass except Exception as exc: + # cancel() makes the call fail, which isn't an error here. + if self._cancelled: + return if not self._resource_exhausted_exception_is_retryable(exc): raise self._reconnect() + def cancel(self): + """Cancel the ``ReadRows`` call. + + Use this to stop reading before the end of the stream. After + ``cancel()``, iterating the stream yields no more messages, and the + stream doesn't reconnect. The client stays open for other calls. + """ + self._cancelled = True + cancel = getattr(self._wrapped, "cancel", None) + if cancel is not None: + cancel() + + def __del__(self): + if not getattr(self, "_cancelled", True): + self.cancel() + def _reconnect(self): """Reconnect to the ReadRows stream using the most recent offset.""" - while True: + while not self._cancelled: try: - self._wrapped = self._client.read_rows( - read_stream=self._name, - offset=self._offset, - **self._read_rows_kwargs, - ) + if self._request_template is None: + wrapped = self._client.read_rows( + read_stream=self._name, + offset=self._offset, + **self._read_rows_kwargs, + ) + else: + # The flattened read_rows() arguments only include + # read_stream and offset, so send a full request copy. + request = gapic_types.ReadRowsRequest(self._request_template) + request.read_stream = self._name + request.offset = self._offset + wrapped = self._client.read_rows( + request=request, + **self._read_rows_kwargs, + ) + # Store the new call before checking for cancel(), so a + # concurrent cancel() either cancels it or is seen here. + self._wrapped = wrapped + if self._cancelled: + self.cancel() + return break except Exception as exc: + if self._cancelled: + return if not self._resource_exhausted_exception_is_retryable(exc): raise diff --git a/packages/google-cloud-bigquery-storage/tests/system/reader/test_reader.py b/packages/google-cloud-bigquery-storage/tests/system/reader/test_reader.py index e9269bda9d96..a9279c89f05c 100644 --- a/packages/google-cloud-bigquery-storage/tests/system/reader/test_reader.py +++ b/packages/google-cloud-bigquery-storage/tests/system/reader/test_reader.py @@ -536,3 +536,67 @@ def test_read_rows_to_dataframe_with_wide_table(client_and_types, project_id): some_rows = next(pages_iter) assert all(len(row["tract_geom"].as_py()) > 0 for row in some_rows) + + +def _arrow_table(responses): + import pyarrow + + schema, batches = None, [] + for response in responses: + if schema is None and response.arrow_schema.serialized_schema: + schema = pyarrow.ipc.read_schema( + pyarrow.py_buffer(response.arrow_schema.serialized_schema) + ) + if response.arrow_record_batch.serialized_record_batch: + batches.append( + pyarrow.ipc.read_record_batch( + pyarrow.py_buffer( + response.arrow_record_batch.serialized_record_batch + ), + schema, + ) + ) + return pyarrow.Table.from_batches(batches, schema=schema) + + +@pytest.fixture(scope="module") +def job_default_stream(bq_client): + job = bq_client.query( + "SELECT x, REPEAT('A', 1000) AS repeated_str " + "FROM UNNEST(GENERATE_ARRAY(1, 20000)) AS x ORDER BY x" + ) + job.result() + return ( + f"projects/{job.project}/locations/{job.location}" + f"/jobs/{job.job_id}/streams/_default" + ) + + +@pytest.mark.parametrize("codec", ["LZ4_FRAME", "ZSTD"]) +def test_read_rows_job_stream_with_arrow_serialization_options( + client_v1, job_default_stream, codec +): + plain = list(client_v1.read_rows(job_default_stream)) + compressed = list( + client_v1.read_rows( + job_default_stream, + arrow_serialization_options={"buffer_compression": codec}, + ) + ) + + def size(responses): + return sum(len(r.arrow_record_batch.serialized_record_batch) for r in responses) + + assert size(compressed) < size(plain) + assert _arrow_table(compressed).equals(_arrow_table(plain)) + + +def test_read_rows_job_stream_cancel(client_v1, job_default_stream): + stream = client_v1.read_rows(job_default_stream) + messages = iter(stream) + first = next(messages) + + stream.cancel() + + assert first.row_count > 0 + assert list(messages) == [] diff --git a/packages/google-cloud-bigquery-storage/tests/unit/test_read_client_v1.py b/packages/google-cloud-bigquery-storage/tests/unit/test_read_client_v1.py index abb88ba36edc..bca9647bdffd 100644 --- a/packages/google-cloud-bigquery-storage/tests/unit/test_read_client_v1.py +++ b/packages/google-cloud-bigquery-storage/tests/unit/test_read_client_v1.py @@ -153,6 +153,58 @@ def test_read_rows(mock_transport, client_under_test): ) +def test_read_rows_with_arrow_serialization_options(mock_transport, client_under_test): + rpc_callable = mock.Mock() + mock_transport._wrapped_methods[mock_transport.read_rows] = rpc_callable + options = {"buffer_compression": "LZ4_FRAME"} + + client_under_test.read_rows( + "teststream", offset=10, arrow_serialization_options=options + ) + + rpc_callable.assert_called_once_with( + types.ReadRowsRequest( + read_stream="teststream", offset=10, arrow_serialization_options=options + ), + metadata=mock.ANY, + retry=mock.ANY, + timeout=mock.ANY, + ) + + +def test_read_rows_reconnect_resends_arrow_serialization_options( + mock_transport, client_under_test +): + def first_stream(): + yield types.ReadRowsResponse(row_count=5) + raise google.api_core.exceptions.ServiceUnavailable("transient disconnect") + + rpc_callable = mock.Mock( + side_effect=[first_stream(), iter([types.ReadRowsResponse(row_count=3)])] + ) + mock_transport._wrapped_methods[mock_transport.read_rows] = rpc_callable + options = {"buffer_compression": "ZSTD"} + + stream = client_under_test.read_rows( + "teststream", offset=10, arrow_serialization_options=options + ) + + assert [m.row_count for m in stream] == [5, 3] + assert rpc_callable.call_args_list == [ + mock.call( + types.ReadRowsRequest( + read_stream="teststream", + offset=offset, + arrow_serialization_options=options, + ), + metadata=mock.ANY, + retry=mock.ANY, + timeout=mock.ANY, + ) + for offset in (10, 15) + ] + + @pytest.mark.parametrize( "module_under_test", ["google.cloud.bigquery_storage_v1", "google.cloud.bigquery_storage_v1beta2"], diff --git a/packages/google-cloud-bigquery-storage/tests/unit/test_reader_v1.py b/packages/google-cloud-bigquery-storage/tests/unit/test_reader_v1.py index d5a83b2a7cf7..ed8220fe177f 100644 --- a/packages/google-cloud-bigquery-storage/tests/unit/test_reader_v1.py +++ b/packages/google-cloud-bigquery-storage/tests/unit/test_reader_v1.py @@ -717,3 +717,151 @@ def test_to_dataframe_emits_pending_deprecation_warning(mut): assert actual_df == mock_df mock_parser.to_dataframe.assert_called_once_with(mock_message, dtypes=None) + + +def test_cancel_stops_iteration_without_reconnecting( + class_under_test, mock_gapic_client +): + first = types.ReadRowsResponse(row_count=1) + + def pages(): + yield first + raise google.api_core.exceptions.ServiceUnavailable("connection closed") + + wrapped = mock.MagicMock() + wrapped.__iter__.return_value = pages() + mock_gapic_client.read_rows.return_value = wrapped + reader = class_under_test(mock_gapic_client, "teststream", 0, {}) + + messages = iter(reader) + assert next(messages) is first + reader.cancel() + + assert list(messages) == [] + wrapped.cancel.assert_called_once_with() + mock_gapic_client.read_rows.assert_called_once_with( + read_stream="teststream", offset=0 + ) + + +def test_cancel_ignores_the_cancelled_call_error(class_under_test, mock_gapic_client): + def pages(): + yield types.ReadRowsResponse(row_count=1) + raise google.api_core.exceptions.Cancelled("call cancelled") + + mock_gapic_client.read_rows.return_value = pages() + reader = class_under_test(mock_gapic_client, "teststream", 0, {}) + + messages = iter(reader) + next(messages) + reader.cancel() + + assert list(messages) == [] + + +def test_cancel_before_iterating_yields_nothing(class_under_test, mock_gapic_client): + reader = class_under_test(mock_gapic_client, "teststream", 0, {}) + + reader.cancel() + + assert list(reader) == [] + mock_gapic_client.read_rows.assert_not_called() + + +def test_cancel_during_reconnect_cancels_new_stream( + class_under_test, mock_gapic_client +): + reader = class_under_test(mock_gapic_client, "teststream", 0, {}) + wrapped = mock.MagicMock() + + def cancel_while_opening(**kwargs): + reader.cancel() + return wrapped + + mock_gapic_client.read_rows.side_effect = cancel_while_opening + + assert list(reader) == [] + wrapped.cancel.assert_called_once_with() + assert reader._wrapped is wrapped + + +def test_cancel_during_reconnect_retry_stops_without_retrying( + class_under_test, mock_gapic_client +): + reader = class_under_test(mock_gapic_client, "teststream", 0, {}) + + def cancel_and_fail(**kwargs): + reader.cancel() + raise google.api_core.exceptions.Cancelled("call cancelled") + + mock_gapic_client.read_rows.side_effect = cancel_and_fail + + assert list(reader) == [] + mock_gapic_client.read_rows.assert_called_once_with( + read_stream="teststream", offset=0 + ) + + +def test_del_cancels_unfinished_stream_once(class_under_test, mock_gapic_client): + wrapped = mock.MagicMock() + mock_gapic_client.read_rows.return_value = wrapped + reader = class_under_test(mock_gapic_client, "teststream", 0, {}) + reader._reconnect() + + del reader + + wrapped.cancel.assert_called_once_with() + + +def test_del_does_not_cancel_twice_or_after_completion( + class_under_test, mock_gapic_client +): + cancelled = mock.MagicMock() + completed = mock.MagicMock() + completed.__iter__.return_value = iter([types.ReadRowsResponse(row_count=1)]) + mock_gapic_client.read_rows.side_effect = [cancelled, completed] + + reader = class_under_test(mock_gapic_client, "teststream", 0, {}) + reader._reconnect() + reader.cancel() + del reader + + finished_reader = class_under_test(mock_gapic_client, "teststream", 0, {}) + assert len(list(finished_reader)) == 1 + del finished_reader + + cancelled.cancel.assert_called_once_with() + completed.cancel.assert_not_called() + + +def test_request_template_not_mutated_across_reconnects( + class_under_test, mock_gapic_client +): + template = types.ReadRowsRequest( + arrow_serialization_options={"buffer_compression": "LZ4_FRAME"} + ) + mock_gapic_client.read_rows.side_effect = [ + _pages_w_unavailable([types.ReadRowsResponse(row_count=4)]), + iter([types.ReadRowsResponse(row_count=2)]), + ] + reader = class_under_test( + mock_gapic_client, + "teststream", + 10, + {}, + request_template=template, + ) + + assert [m.row_count for m in reader] == [4, 2] + assert template.read_stream == "" + assert template.offset == 0 + assert mock_gapic_client.read_rows.call_args_list == [ + mock.call( + request=types.ReadRowsRequest( + read_stream="teststream", + offset=offset, + arrow_serialization_options={"buffer_compression": "LZ4_FRAME"}, + ) + ) + for offset in (10, 14) + ]