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 ade01abb67b8..5a8bf4271458 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 @@ -83,6 +83,13 @@ class ReadRowsStream(object): :func:`~google.cloud.bigquery_storage_v1.reader.ReadRowsStream.to_dataframe()` method to parse all messages into a :class:`pandas.DataFrame`. + Use the stream as a context manager to cancel the ``ReadRows`` call if + reading stops before the end of the stream:: + + with client.read_rows(stream_name) as stream: + for message in stream: + ... + This object should not be created directly, but is returned by other methods in this library. """ @@ -147,6 +154,7 @@ def __init__( self._retry_delay_callback = retry_delay_callback self._request_template = request_template self._cancelled = False + self._finished = False self._wrapped = None def __iter__(self): @@ -172,6 +180,7 @@ def __iter__(self): yield message self._wrapped = None + self._finished = True return # Made it through the whole stream. except google.api_core.exceptions.InternalServerError as exc: resumable_error = any( @@ -205,7 +214,16 @@ def cancel(self): cancel() def __del__(self): - if not getattr(self, "_cancelled", True): + if not getattr(self, "_cancelled", True) and not getattr( + self, "_finished", True + ): + self.cancel() + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_value, traceback): + if not self._cancelled and not self._finished: self.cancel() def _reconnect(self): 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 ed8220fe177f..48b885df4a78 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 @@ -834,6 +834,63 @@ def test_del_does_not_cancel_twice_or_after_completion( completed.cancel.assert_not_called() +def test_del_does_not_call_cancel_after_completion(class_under_test, mock_gapic_client): + mock_gapic_client.read_rows.return_value = iter( + [types.ReadRowsResponse(row_count=1)] + ) + reader = class_under_test(mock_gapic_client, "teststream", 0, {}) + assert len(list(reader)) == 1 + + with mock.patch.object(class_under_test, "cancel", autospec=True) as cancel: + del reader + + cancel.assert_not_called() + + +def test_context_manager_cancels_unfinished_stream(class_under_test, mock_gapic_client): + wrapped = mock.MagicMock() + wrapped.__iter__.return_value = iter( + [types.ReadRowsResponse(row_count=1), types.ReadRowsResponse(row_count=1)] + ) + mock_gapic_client.read_rows.return_value = wrapped + reader = class_under_test(mock_gapic_client, "teststream", 0, {}) + + with reader as stream: + assert stream is reader + next(iter(stream)) + + wrapped.cancel.assert_called_once_with() + + +def test_context_manager_does_not_cancel_finished_stream( + class_under_test, mock_gapic_client +): + mock_gapic_client.read_rows.return_value = iter( + [types.ReadRowsResponse(row_count=1)] + ) + + with mock.patch.object(class_under_test, "cancel", autospec=True) as cancel: + with class_under_test(mock_gapic_client, "teststream", 0, {}) as stream: + assert len(list(stream)) == 1 + + cancel.assert_not_called() + + +def test_context_manager_cancels_and_reraises_on_error( + 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, {}) + + with pytest.raises(ValueError, match="stop reading"): + with reader: + reader._reconnect() + raise ValueError("stop reading") + + wrapped.cancel.assert_called_once_with() + + def test_request_template_not_mutated_across_reconnects( class_under_test, mock_gapic_client ):