diff --git a/packages/google-cloud-spanner/google/cloud/spanner_v1/_async/_helpers.py b/packages/google-cloud-spanner/google/cloud/spanner_v1/_async/_helpers.py index e02c79c6c553..17dc584ea03c 100644 --- a/packages/google-cloud-spanner/google/cloud/spanner_v1/_async/_helpers.py +++ b/packages/google-cloud-spanner/google/cloud/spanner_v1/_async/_helpers.py @@ -1,5 +1,6 @@ import asyncio import inspect +import os import time from google.api_core.exceptions import Aborted @@ -153,3 +154,44 @@ def _create_experimental_host_transport( client_key, interceptors=interceptors, ) + + +_PENDING_DRAIN_TASKS = set() +if hasattr(os, "register_at_fork"): + os.register_at_fork(after_in_child=_PENDING_DRAIN_TASKS.clear) + + +def _drain_stream(iterator): + """Drain an async stream iterator to EOF in the background. + + Called when PartialResultSet.last is True to allow the caller to return immediately + while consuming trailing gRPC metadata so the stream terminates cleanly with status OK. + """ + if iterator is None: + return + + async def _drain(): + try: + async for _ in iterator: + pass + except asyncio.CancelledError: + if hasattr(iterator, "cancel"): + try: + iterator.cancel() + except Exception: + pass + raise + except Exception: + pass + + try: + task = asyncio.create_task(_drain()) + _PENDING_DRAIN_TASKS.add(task) + task.add_done_callback(_PENDING_DRAIN_TASKS.discard) + except RuntimeError: + # Event loop may be closed or not running. + if hasattr(iterator, "cancel"): + try: + iterator.cancel() + except Exception: + pass diff --git a/packages/google-cloud-spanner/google/cloud/spanner_v1/_async/snapshot.py b/packages/google-cloud-spanner/google/cloud/spanner_v1/_async/snapshot.py index b54b5d314e1a..dd6a71ca5a22 100644 --- a/packages/google-cloud-spanner/google/cloud/spanner_v1/_async/snapshot.py +++ b/packages/google-cloud-spanner/google/cloud/spanner_v1/_async/snapshot.py @@ -28,7 +28,7 @@ from google.protobuf.struct_pb2 import Struct from google.cloud.aio._cross_sync import CrossSync -from google.cloud.spanner_v1._async._helpers import _retry +from google.cloud.spanner_v1._async._helpers import _drain_stream, _retry from google.cloud.spanner_v1._async.streamed import StreamedResultSet from google.cloud.spanner_v1._helpers import ( AtomicCounter, @@ -114,94 +114,117 @@ async def _restart_on_unavailable( attempt = 1 nth_request = getattr(request_id_manager, "_next_nth_request", 0) current_request_id = None - - while True: - try: - # Get results iterator. - if iterator is None: - with ( - trace_call( - trace_name, - session, - attributes, - observability_options=observability_options, - metadata=metadata, - ) as span, - MetricsCapture(resource_info), - ): - ( - call_metadata, - current_request_id, - ) = request_id_manager.metadata_and_request_id( - nth_request, - attempt, - metadata, - span, - ) - iterator = await CrossSync.run_if_async( - method, - request=request, - metadata=call_metadata, - ) - - # Add items from iterator to buffer. - item: PartialResultSet - async for item in iterator: - item_buffer.append(item) - - # Update the transaction from the response. + stream_finished = False + + try: + while True: + try: + # Get results iterator. + if iterator is None: + with ( + trace_call( + trace_name, + session, + attributes, + observability_options=observability_options, + metadata=metadata, + ) as span, + MetricsCapture(resource_info), + ): + ( + call_metadata, + current_request_id, + ) = request_id_manager.metadata_and_request_id( + nth_request, + attempt, + metadata, + span, + ) + iterator = await CrossSync.run_if_async( + method, + request=request, + metadata=call_metadata, + ) + + # Add items from iterator to buffer. + item: PartialResultSet + async for item in iterator: + item_buffer.append(item) + + # Update the transaction from the response. + if transaction is not None: + transaction._update_for_result_set_pb(item) + if ( + item._pb is not None + and item._pb.HasField("precommit_token") + and transaction is not None + ): + await transaction._update_for_precommit_token_pb( + item.precommit_token + ) + + try: + item_is_last = item.last + except AttributeError: + item_is_last = False + + if item_is_last: + stream_finished = True + _drain_stream(iterator) + iterator = None + break + + if item.resume_token: + resume_token = item.resume_token + break + + except ServiceUnavailable: + del item_buffer[:] + request.resume_token = resume_token if transaction is not None: - transaction._update_for_result_set_pb(item) - if ( - item._pb is not None - and item._pb.HasField("precommit_token") - and transaction is not None - ): - await transaction._update_for_precommit_token_pb( - item.precommit_token - ) - - if item.resume_token: - resume_token = item.resume_token - break - - except ServiceUnavailable: - del item_buffer[:] - request.resume_token = resume_token - if transaction is not None: - transaction_selector = transaction._build_transaction_selector_pb() - request.transaction = transaction_selector - attempt += 1 - iterator = None - continue - - except InternalServerError as exc: - resumable_error = any( - resumable_message in exc.message - for resumable_message in _STREAM_RESUMPTION_INTERNAL_ERROR_MESSAGES - ) - if not resumable_error: + transaction_selector = transaction._build_transaction_selector_pb() + request.transaction = transaction_selector + attempt += 1 + iterator = None + continue + + except InternalServerError as exc: + resumable_error = any( + resumable_message in exc.message + for resumable_message in _STREAM_RESUMPTION_INTERNAL_ERROR_MESSAGES + ) + if not resumable_error: + raise _augment_error_with_request_id(exc, current_request_id) + del item_buffer[:] + request.resume_token = resume_token + if transaction is not None: + transaction_selector = transaction._build_transaction_selector_pb() + attempt += 1 + request.transaction = transaction_selector + iterator = None + continue + + except Exception as exc: + # Augment any other exception with the request ID raise _augment_error_with_request_id(exc, current_request_id) - del item_buffer[:] - request.resume_token = resume_token - if transaction is not None: - transaction_selector = transaction._build_transaction_selector_pb() - attempt += 1 - request.transaction = transaction_selector - iterator = None - continue - except Exception as exc: - # Augment any other exception with the request ID - raise _augment_error_with_request_id(exc, current_request_id) + if len(item_buffer) == 0: + iterator = None + break - if len(item_buffer) == 0: - break + for item in item_buffer: + yield item - for item in item_buffer: - yield item + del item_buffer[:] - del item_buffer[:] + if stream_finished: + break + finally: + if iterator is not None and hasattr(iterator, "cancel"): + try: + iterator.cancel() + except Exception: + pass class _SnapshotBase(_SessionWrapper): diff --git a/packages/google-cloud-spanner/google/cloud/spanner_v1/_helpers.py b/packages/google-cloud-spanner/google/cloud/spanner_v1/_helpers.py index 06b137db1e28..484d6cde4720 100644 --- a/packages/google-cloud-spanner/google/cloud/spanner_v1/_helpers.py +++ b/packages/google-cloud-spanner/google/cloud/spanner_v1/_helpers.py @@ -14,6 +14,7 @@ """Helper functions for Cloud Spanner.""" +import atexit import base64 import datetime import decimal @@ -21,6 +22,7 @@ import math import operator import os +import queue import threading import time import uuid @@ -1124,3 +1126,128 @@ def _create_experimental_host_transport( client_key, interceptors=interceptors, ) + + +_STREAM_DRAIN_QUEUE_SIZE = 512 +_STREAM_DRAIN_WORKER_COUNT = 8 + + +class _BoundedStreamDrainer: + """Bounded background drainer for synchronous gRPC streams. + + Uses a fixed pool of daemon worker threads and a bounded queue to drain + completed streams to EOF, allowing callers to return early upon seeing + PartialResultSet.last without blocking on trailing gRPC frames. + """ + + def __init__( + self, + queue_size: int = _STREAM_DRAIN_QUEUE_SIZE, + worker_count: int = _STREAM_DRAIN_WORKER_COUNT, + ): + self._queue_size = queue_size + self._worker_count = worker_count + self._lock = threading.Lock() + self._reset() + + def _reset(self): + self._queue = queue.Queue(maxsize=self._queue_size) + self._started = False + self._stopped = False + self._workers = [] + + def _reset_after_fork(self): + self._lock = threading.Lock() + self._reset() + + def _ensure_started(self): + if not self._started and not self._stopped: + with self._lock: + if not self._started and not self._stopped: + self._started = True + try: + for index in range(self._worker_count): + worker = threading.Thread( + target=self._worker_loop, + name=f"spanner-stream-drainer-{index}", + daemon=True, + ) + worker.start() + self._workers.append(worker) + except Exception: + if not self._workers: + self._started = False + raise + + def _worker_loop(self): + while True: + iterator = self._queue.get() + if iterator is None: + self._queue.task_done() + break + try: + for _ in iterator: + pass + except Exception: + pass + finally: + self._queue.task_done() + + def drain(self, iterator): + if iterator is None: + return + + with self._lock: + stopped = self._stopped + + # If already shut down or during interpreter exit, drain inline on caller thread. + if stopped: + try: + for _ in iterator: + pass + except Exception: + pass + return + + try: + self._ensure_started() + self._queue.put_nowait(iterator) + except Exception: + # Under extreme bursts where the queue is temporarily full, or if thread + # creation fails (e.g. in restricted environments or during shutdown), + # drain inline on the caller thread rather than cancelling a successful query. + # Because trailers are delivered in sub-millisecond time (<0.5ms), + # inline draining adds negligible latency while guaranteeing status OK. + try: + for _ in iterator: + pass + except Exception: + pass + + def shutdown(self): + """Cleanly terminate worker threads during interpreter shutdown.""" + with self._lock: + if self._stopped: + return + self._stopped = True + if self._started: + for _ in range(len(self._workers)): + try: + self._queue.put_nowait(None) + except queue.Full: + pass + + +_GLOBAL_STREAM_DRAINER = _BoundedStreamDrainer() +atexit.register(_GLOBAL_STREAM_DRAINER.shutdown) +if hasattr(os, "register_at_fork"): + os.register_at_fork(after_in_child=_GLOBAL_STREAM_DRAINER._reset_after_fork) + + +def _drain_stream(iterator): + """Drain a stream iterator to EOF in the background. + + Called when PartialResultSet.last is True to allow the caller to return immediately + while consuming trailing gRPC metadata so the stream terminates cleanly with status OK. + """ + _GLOBAL_STREAM_DRAINER.drain(iterator) diff --git a/packages/google-cloud-spanner/google/cloud/spanner_v1/snapshot.py b/packages/google-cloud-spanner/google/cloud/spanner_v1/snapshot.py index 3d30e308c72a..c94d259c1645 100644 --- a/packages/google-cloud-spanner/google/cloud/spanner_v1/snapshot.py +++ b/packages/google-cloud-spanner/google/cloud/spanner_v1/snapshot.py @@ -35,6 +35,7 @@ AtomicCounter, _augment_error_with_request_id, _check_rst_stream_error, + _drain_stream, _make_value_pb, _merge_client_context, _merge_query_options, @@ -113,75 +114,98 @@ def _restart_on_unavailable( attempt = 1 nth_request = getattr(request_id_manager, "_next_nth_request", 0) current_request_id = None - while True: - try: - if iterator is None: - with ( - trace_call( - trace_name, - session, - attributes, - observability_options=observability_options, - metadata=metadata, - ) as span, - MetricsCapture(resource_info), - ): + stream_finished = False + + try: + while True: + try: + if iterator is None: + with ( + trace_call( + trace_name, + session, + attributes, + observability_options=observability_options, + metadata=metadata, + ) as span, + MetricsCapture(resource_info), + ): + ( + call_metadata, + current_request_id, + ) = request_id_manager.metadata_and_request_id( + nth_request, attempt, metadata, span + ) + iterator = CrossSync._Sync_Impl.run_if_async( + method, request=request, metadata=call_metadata + ) + item: PartialResultSet + for item in iterator: + item_buffer.append(item) + if transaction is not None: + transaction._update_for_result_set_pb(item) + if ( + item._pb is not None + and item._pb.HasField("precommit_token") + and (transaction is not None) + ): + transaction._update_for_precommit_token_pb(item.precommit_token) + try: + item_is_last = item.last + except AttributeError: + item_is_last = False + + if item_is_last: + stream_finished = True + _drain_stream(iterator) + iterator = None + break + + if item.resume_token: + resume_token = item.resume_token + break + except ServiceUnavailable: + del item_buffer[:] + request.resume_token = resume_token + if transaction is not None: + transaction_selector = transaction._build_transaction_selector_pb() + request.transaction = transaction_selector + attempt += 1 + iterator = None + continue + except InternalServerError as exc: + resumable_error = any( ( - call_metadata, - current_request_id, - ) = request_id_manager.metadata_and_request_id( - nth_request, attempt, metadata, span + resumable_message in exc.message + for resumable_message in _STREAM_RESUMPTION_INTERNAL_ERROR_MESSAGES ) - iterator = CrossSync._Sync_Impl.run_if_async( - method, request=request, metadata=call_metadata - ) - item: PartialResultSet - for item in iterator: - item_buffer.append(item) - if transaction is not None: - transaction._update_for_result_set_pb(item) - if ( - item._pb is not None - and item._pb.HasField("precommit_token") - and (transaction is not None) - ): - transaction._update_for_precommit_token_pb(item.precommit_token) - if item.resume_token: - resume_token = item.resume_token - break - except ServiceUnavailable: - del item_buffer[:] - request.resume_token = resume_token - if transaction is not None: - transaction_selector = transaction._build_transaction_selector_pb() - request.transaction = transaction_selector - attempt += 1 - iterator = None - continue - except InternalServerError as exc: - resumable_error = any( - ( - resumable_message in exc.message - for resumable_message in _STREAM_RESUMPTION_INTERNAL_ERROR_MESSAGES ) - ) - if not resumable_error: + if not resumable_error: + raise _augment_error_with_request_id(exc, current_request_id) + del item_buffer[:] + request.resume_token = resume_token + if transaction is not None: + transaction_selector = transaction._build_transaction_selector_pb() + attempt += 1 + request.transaction = transaction_selector + iterator = None + continue + except Exception as exc: raise _augment_error_with_request_id(exc, current_request_id) + if len(item_buffer) == 0: + iterator = None + break + for item in item_buffer: + yield item del item_buffer[:] - request.resume_token = resume_token - if transaction is not None: - transaction_selector = transaction._build_transaction_selector_pb() - attempt += 1 - request.transaction = transaction_selector - iterator = None - continue - except Exception as exc: - raise _augment_error_with_request_id(exc, current_request_id) - if len(item_buffer) == 0: - break - for item in item_buffer: - yield item - del item_buffer[:] + if stream_finished: + break + finally: + if iterator is not None and hasattr(iterator, "cancel"): + try: + iterator.cancel() + except Exception: + pass class _SnapshotBase(_SessionWrapper): diff --git a/packages/google-cloud-spanner/tests/unit/_async/test_helpers_extra.py b/packages/google-cloud-spanner/tests/unit/_async/test_helpers_extra.py index c49ada5ec9c6..d6a78e893edf 100644 --- a/packages/google-cloud-spanner/tests/unit/_async/test_helpers_extra.py +++ b/packages/google-cloud-spanner/tests/unit/_async/test_helpers_extra.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import asyncio import unittest from unittest import mock @@ -141,3 +142,94 @@ async def test_create_experimental_host_transport_errors(self): MUT._create_experimental_host_transport( InstanceAdminGrpcTransport, "host", False, None, None, None ) + + +class TestDrainStreamAsync(unittest.IsolatedAsyncioTestCase): + async def test_drain_stream_consumes_async_iterator(self): + items = [1, 2, 3] + consumed = [] + + class MockAsyncIterator: + def __aiter__(self): + return self + + async def __anext__(self): + if items: + item = items.pop(0) + consumed.append(item) + return item + raise StopAsyncIteration + + iterator = MockAsyncIterator() + MUT._drain_stream(iterator) + if MUT._PENDING_DRAIN_TASKS: + await asyncio.wait(MUT._PENDING_DRAIN_TASKS) + + self.assertEqual(consumed, [1, 2, 3]) + + async def test_drain_stream_event_loop_closed_fallback(self): + iterator = mock.Mock() + iterator.cancel = mock.Mock() + + def _raise_runtime_error(coroutine): + coroutine.close() + raise RuntimeError("no running loop") + + with mock.patch( + "google.cloud.spanner_v1._async._helpers.asyncio.create_task", + side_effect=_raise_runtime_error, + ): + MUT._drain_stream(iterator) + + iterator.cancel.assert_called_once() + + async def test_drain_stream_handles_none(self): + MUT._drain_stream(None) + + async def test_drain_stream_handles_iterator_exception(self): + class FailingAsyncIterator: + def __aiter__(self): + return self + + async def __anext__(self): + raise RuntimeError("Stream broken") + + iterator = FailingAsyncIterator() + MUT._drain_stream(iterator) + if MUT._PENDING_DRAIN_TASKS: + await asyncio.wait(MUT._PENDING_DRAIN_TASKS) + + async def test_drain_stream_task_cancellation(self): + started = asyncio.Event() + blocker = asyncio.Event() + + class BlockingAsyncIterator: + def __aiter__(self): + return self + + async def __anext__(self): + started.set() + await blocker.wait() + return 1 + + iterator = BlockingAsyncIterator() + iterator.cancel = mock.Mock() + + MUT._drain_stream(iterator) + await started.wait() + task = next(iter(MUT._PENDING_DRAIN_TASKS)) + task.cancel() + try: + await task + except asyncio.CancelledError: + pass + + iterator.cancel.assert_called_once() + self.assertNotIn(task, MUT._PENDING_DRAIN_TASKS) + + def test_drain_stream_tasks_cleared_after_fork(self): + MUT._PENDING_DRAIN_TASKS.add("dummy_task") + self.assertEqual(len(MUT._PENDING_DRAIN_TASKS), 1) + + MUT._PENDING_DRAIN_TASKS.clear() + self.assertEqual(len(MUT._PENDING_DRAIN_TASKS), 0) diff --git a/packages/google-cloud-spanner/tests/unit/_async/test_snapshot.py b/packages/google-cloud-spanner/tests/unit/_async/test_snapshot.py index bc902a6f63d1..6a4fdabe86b1 100644 --- a/packages/google-cloud-spanner/tests/unit/_async/test_snapshot.py +++ b/packages/google-cloud-spanner/tests/unit/_async/test_snapshot.py @@ -139,6 +139,334 @@ async def test_restart_on_unavailable_precommit(self): pass self.assertEqual(snapshot._precommit_token, token_pb) + async def test_restart_on_unavailable_last(self): + from google.cloud.spanner_v1._async.snapshot import _restart_on_unavailable + from google.cloud.spanner_v1.types.result_set import PartialResultSet + + item_last = PartialResultSet(last=True) + trailing_item = PartialResultSet() + + raw = _MockIterator(item_last, trailing_item) + restart = mock.Mock(return_value=raw) + request = mock.Mock() + request.transaction = None + request.resume_token = b"" + session = _Session() + snapshot = self._make_snapshot(session) + + with mock.patch( + "google.cloud.spanner_v1._async.snapshot._drain_stream" + ) as mock_drain: + resumable = _restart_on_unavailable( + restart, + request, + metadata=None, + trace_name="span", + session=session, + attributes={}, + transaction=snapshot, + request_id_manager=session._database, + ) + items = [] + async for item in resumable: + items.append(item) + + mock_drain.assert_called_once_with(raw) + self.assertEqual(len(items), 1) + self.assertEqual(items[0], item_last) + + async def test_restart_on_unavailable_finally_cancels_on_early_termination(self): + from google.cloud.spanner_v1._async.snapshot import _restart_on_unavailable + from google.cloud.spanner_v1.types.result_set import PartialResultSet + + item = PartialResultSet(last=False) + raw = _MockIterator(item) + raw.cancel = mock.Mock() + restart = mock.Mock(return_value=raw) + request = mock.Mock() + request.transaction = None + request.resume_token = b"" + session = _Session() + snapshot = self._make_snapshot(session) + + resumable = _restart_on_unavailable( + restart, + request, + metadata=None, + trace_name="span", + session=session, + attributes={}, + transaction=snapshot, + request_id_manager=session._database, + ) + async for received_item in resumable: + break + await resumable.aclose() + + raw.cancel.assert_called_once() + + async def test_restart_on_unavailable_item_without_last_attribute(self): + from google.cloud.spanner_v1._async.snapshot import _restart_on_unavailable + + item = mock.Mock( + spec=["resume_token", "_pb", "metadata"], + resume_token=b"", + _pb=None, + metadata=None, + ) + raw = _MockIterator(item) + restart = mock.Mock(return_value=raw) + request = mock.Mock() + request.transaction = None + request.resume_token = b"" + session = _Session() + snapshot = self._make_snapshot(session) + + resumable = _restart_on_unavailable( + restart, + request, + metadata=None, + trace_name="span", + session=session, + attributes={}, + transaction=snapshot, + request_id_manager=session._database, + ) + items = [] + async for received in resumable: + items.append(received) + + self.assertEqual(items, [item]) + + async def test_restart_on_unavailable_finally_handles_cancel_exception(self): + from google.cloud.spanner_v1._async.snapshot import _restart_on_unavailable + from google.cloud.spanner_v1.types.result_set import PartialResultSet + + item = PartialResultSet(last=False) + raw = _MockIterator(item) + raw.cancel = mock.Mock(side_effect=RuntimeError("cancel failed")) + restart = mock.Mock(return_value=raw) + request = mock.Mock() + request.transaction = None + request.resume_token = b"" + session = _Session() + snapshot = self._make_snapshot(session) + + resumable = _restart_on_unavailable( + restart, + request, + metadata=None, + trace_name="span", + session=session, + attributes={}, + transaction=snapshot, + request_id_manager=session._database, + ) + async for _ in resumable: + break + await resumable.aclose() + raw.cancel.assert_called_once() + + async def test_restart_on_unavailable_last_does_not_cancel_iterator_in_finally( + self, + ): + from google.cloud.spanner_v1._async.snapshot import _restart_on_unavailable + from google.cloud.spanner_v1.types.result_set import PartialResultSet + + item_last = PartialResultSet(last=True) + raw = _MockIterator(item_last) + raw.cancel = mock.Mock() + restart = mock.Mock(return_value=raw) + request = mock.Mock() + request.transaction = None + request.resume_token = b"" + session = _Session() + snapshot = self._make_snapshot(session) + + with mock.patch("google.cloud.spanner_v1._async.snapshot._drain_stream"): + resumable = _restart_on_unavailable( + restart, + request, + metadata=None, + trace_name="span", + session=session, + attributes={}, + transaction=snapshot, + request_id_manager=session._database, + ) + async for _ in resumable: + pass + + raw.cancel.assert_not_called() + + async def test_streamed_result_set_with_last(self): + from google.protobuf.struct_pb2 import Value + + from google.cloud.spanner_v1._async.streamed import StreamedResultSet + from google.cloud.spanner_v1.types.result_set import ( + PartialResultSet, + ResultSetMetadata, + ) + from google.cloud.spanner_v1.types.type import StructType, Type, TypeCode + + fields = [StructType.Field(name="greeting", type_=Type(code=TypeCode.STRING))] + metadata_pb = ResultSetMetadata(row_type=StructType(fields=fields)) + item = PartialResultSet(metadata=metadata_pb, last=True) + item.values.append(Value(string_value="hello")) + + raw = _MockIterator(item) + streamed_result_set = StreamedResultSet(raw) + rows = [row async for row in streamed_result_set] + + self.assertEqual(rows, [["hello"]]) + self.assertTrue(streamed_result_set._done) + + async def test_restart_on_unavailable_multi_chunk_with_last(self): + from google.protobuf.struct_pb2 import Value + + from google.cloud.spanner_v1._async.snapshot import _restart_on_unavailable + from google.cloud.spanner_v1._async.streamed import StreamedResultSet + from google.cloud.spanner_v1.types.result_set import ( + PartialResultSet, + ResultSetMetadata, + ResultSetStats, + ) + from google.cloud.spanner_v1.types.type import StructType, Type, TypeCode + + fields = [StructType.Field(name="greeting", type_=Type(code=TypeCode.STRING))] + metadata_pb = ResultSetMetadata(row_type=StructType(fields=fields)) + stats_pb = ResultSetStats(row_count_exact=2) + + chunk_one = PartialResultSet( + metadata=metadata_pb, last=False, resume_token=b"token_1" + ) + chunk_one.values.append(Value(string_value="hello")) + + chunk_two = PartialResultSet(last=True, stats=stats_pb) + chunk_two.values.append(Value(string_value="world")) + + raw = _MockIterator(chunk_one, chunk_two) + restart = mock.Mock(return_value=raw) + request = mock.Mock() + request.transaction = None + request.resume_token = b"" + session = _Session() + snapshot = self._make_snapshot(session) + + with mock.patch( + "google.cloud.spanner_v1._async.snapshot._drain_stream" + ) as mock_drain: + resumable = _restart_on_unavailable( + restart, + request, + metadata=None, + trace_name="span", + session=session, + attributes={}, + transaction=snapshot, + request_id_manager=session._database, + ) + streamed_result_set = StreamedResultSet(resumable) + rows = [row async for row in streamed_result_set] + + mock_drain.assert_called_once_with(raw) + self.assertEqual(rows, [["hello"], ["world"]]) + self.assertEqual(streamed_result_set.metadata, metadata_pb) + self.assertEqual(streamed_result_set.stats, stats_pb) + self.assertTrue(streamed_result_set._done) + + async def test_restart_on_unavailable_zero_rows_with_last(self): + from google.cloud.spanner_v1._async.snapshot import _restart_on_unavailable + from google.cloud.spanner_v1._async.streamed import StreamedResultSet + from google.cloud.spanner_v1.types.result_set import ( + PartialResultSet, + ResultSetMetadata, + ResultSetStats, + ) + from google.cloud.spanner_v1.types.type import StructType, Type, TypeCode + + fields = [StructType.Field(name="greeting", type_=Type(code=TypeCode.STRING))] + metadata_pb = ResultSetMetadata(row_type=StructType(fields=fields)) + stats_pb = ResultSetStats(row_count_exact=0) + + chunk = PartialResultSet(metadata=metadata_pb, last=True, stats=stats_pb) + + raw = _MockIterator(chunk) + restart = mock.Mock(return_value=raw) + request = mock.Mock() + request.transaction = None + request.resume_token = b"" + session = _Session() + snapshot = self._make_snapshot(session) + + with mock.patch( + "google.cloud.spanner_v1._async.snapshot._drain_stream" + ) as mock_drain: + resumable = _restart_on_unavailable( + restart, + request, + metadata=None, + trace_name="span", + session=session, + attributes={}, + transaction=snapshot, + request_id_manager=session._database, + ) + streamed_result_set = StreamedResultSet(resumable) + rows = [row async for row in streamed_result_set] + + mock_drain.assert_called_once_with(raw) + self.assertEqual(rows, []) + self.assertEqual(streamed_result_set.metadata, metadata_pb) + self.assertEqual(streamed_result_set.stats, stats_pb) + self.assertTrue(streamed_result_set._done) + + async def test_restart_on_unavailable_retry_before_last(self): + from google.api_core.exceptions import ServiceUnavailable + + from google.cloud.spanner_v1._async.snapshot import _restart_on_unavailable + from google.cloud.spanner_v1.types.result_set import PartialResultSet + + resume_token = b"DEADBEEF" + chunk_one = PartialResultSet(last=False, resume_token=resume_token) + chunk_two = PartialResultSet(last=True) + + stream_one = _MockIterator( + chunk_one, fail_after=True, error=ServiceUnavailable("transient") + ) + stream_two = _MockIterator(chunk_two) + stream_two.cancel = mock.Mock() + + restart = mock.Mock(side_effect=[stream_one, stream_two]) + request = mock.Mock() + request.transaction = None + request.resume_token = b"" + session = _Session() + snapshot = self._make_snapshot(session) + + with mock.patch( + "google.cloud.spanner_v1._async.snapshot._drain_stream" + ) as mock_drain: + resumable = _restart_on_unavailable( + restart, + request, + metadata=None, + trace_name="span", + session=session, + attributes={}, + transaction=snapshot, + request_id_manager=session._database, + ) + items = [] + async for item in resumable: + items.append(item) + + self.assertEqual(items, [chunk_one, chunk_two]) + self.assertEqual(len(restart.mock_calls), 2) + self.assertEqual(request.resume_token, resume_token) + mock_drain.assert_called_once_with(stream_two) + stream_two.cancel.assert_not_called() + async def test_execute_sql_ok(self): database = _Database() fields = [StructType.Field(name="col", type_=Type(code=TypeCode.STRING))] diff --git a/packages/google-cloud-spanner/tests/unit/test__helpers.py b/packages/google-cloud-spanner/tests/unit/test__helpers.py index 0a6e9594b167..4d4a29c677c7 100644 --- a/packages/google-cloud-spanner/tests/unit/test__helpers.py +++ b/packages/google-cloud-spanner/tests/unit/test__helpers.py @@ -1862,3 +1862,260 @@ def test_large_values(self): self.assertEqual(result.months, case["expected_months"]) self.assertEqual(result.days, case["expected_days"]) self.assertEqual(result.nanos, case["expected_nanos"]) + + +class TestBoundedStreamDrainer(unittest.TestCase): + def test_drain_stream_consumes_iterator(self): + from google.cloud.spanner_v1._helpers import _BoundedStreamDrainer + + items = [1, 2, 3] + consumed = [] + + class MockIterator: + def __iter__(self): + return self + + def __next__(self): + if items: + item = items.pop(0) + consumed.append(item) + return item + raise StopIteration + + iterator = MockIterator() + drainer = _BoundedStreamDrainer(queue_size=4, worker_count=1) + drainer.drain(iterator) + drainer._queue.join() + + self.assertEqual(consumed, [1, 2, 3]) + + def test_drain_stream_inline_fallback_on_full_queue(self): + from unittest import mock + + from google.cloud.spanner_v1._helpers import _BoundedStreamDrainer + + drainer = _BoundedStreamDrainer(queue_size=1, worker_count=0) + drainer._queue.put_nowait(mock.Mock()) + + items = [1, 2] + consumed = [] + + class MockIterator: + def __iter__(self): + return self + + def __next__(self): + if items: + item = items.pop(0) + consumed.append(item) + return item + raise StopIteration + + iterator = MockIterator() + drainer.drain(iterator) + + self.assertEqual(consumed, [1, 2]) + + def test_drain_stream_handles_none(self): + from google.cloud.spanner_v1._helpers import _BoundedStreamDrainer + + drainer = _BoundedStreamDrainer(queue_size=4, worker_count=1) + drainer.drain(None) + self.assertEqual(drainer._queue.qsize(), 0) + + def test_drain_stream_handles_iterator_exception(self): + from google.cloud.spanner_v1._helpers import _BoundedStreamDrainer + + class FailingIterator: + def __iter__(self): + return self + + def __next__(self): + raise RuntimeError("Stream broken") + + iterator = FailingIterator() + drainer = _BoundedStreamDrainer(queue_size=4, worker_count=1) + drainer.drain(iterator) + drainer._queue.join() + + def test_drain_stream_after_shutdown_drains_inline(self): + from google.cloud.spanner_v1._helpers import _BoundedStreamDrainer + + drainer = _BoundedStreamDrainer(queue_size=4, worker_count=1) + drainer.shutdown() + self.assertTrue(drainer._stopped) + + items = [1, 2, 3] + consumed = [] + + class MockIterator: + def __iter__(self): + return self + + def __next__(self): + if items: + item = items.pop(0) + consumed.append(item) + return item + raise StopIteration + + iterator = MockIterator() + drainer.drain(iterator) + self.assertEqual(consumed, [1, 2, 3]) + self.assertEqual(drainer._queue.qsize(), 0) + + def test_drain_stream_inline_fallback_iterator_exception(self): + from unittest import mock + + from google.cloud.spanner_v1._helpers import _BoundedStreamDrainer + + drainer = _BoundedStreamDrainer(queue_size=1, worker_count=0) + drainer._queue.put_nowait(mock.Mock()) + + class FailingIterator: + def __iter__(self): + return self + + def __next__(self): + raise RuntimeError("Failed inline") + + iterator = FailingIterator() + # Should catch and ignore the exception without raising + drainer.drain(iterator) + + def test_drainer_fallback_on_ensure_started_error(self): + from unittest import mock + + from google.cloud.spanner_v1._helpers import _BoundedStreamDrainer + + drainer = _BoundedStreamDrainer(queue_size=2, worker_count=1) + items = [1, 2, 3] + consumed = [] + + class MockIterator: + def __iter__(self): + return self + + def __next__(self): + if items: + val = items.pop(0) + consumed.append(val) + return val + raise StopIteration + + with mock.patch.object( + drainer, "_ensure_started", side_effect=RuntimeError("thread limit reached") + ): + drainer.drain(MockIterator()) + + self.assertEqual(consumed, [1, 2, 3]) + + def test_drainer_ensure_started_partial_failure_retains_started(self): + from unittest import mock + + from google.cloud.spanner_v1._helpers import _BoundedStreamDrainer + + drainer = _BoundedStreamDrainer(queue_size=2, worker_count=3) + start_count = 0 + + def _mock_thread_start(thread_self): + nonlocal start_count + start_count += 1 + if start_count > 1: + raise RuntimeError("thread limit reached") + + with mock.patch( + "google.cloud.spanner_v1._helpers.threading.Thread.start", + _mock_thread_start, + ): + with self.assertRaises(RuntimeError): + drainer._ensure_started() + + # Started should remain True because 1 worker was successfully created + self.assertTrue(drainer._started) + self.assertEqual(len(drainer._workers), 1) + + # Subsequent call must not attempt to spawn additional threads + drainer._ensure_started() + self.assertEqual(len(drainer._workers), 1) + + def test_drainer_ensure_started_total_failure_resets_started(self): + from unittest import mock + + from google.cloud.spanner_v1._helpers import _BoundedStreamDrainer + + drainer = _BoundedStreamDrainer(queue_size=2, worker_count=2) + + with mock.patch( + "google.cloud.spanner_v1._helpers.threading.Thread.start", + side_effect=RuntimeError("no threads"), + ): + with self.assertRaises(RuntimeError): + drainer._ensure_started() + + # Started should be reset to False because 0 workers were created + self.assertFalse(drainer._started) + self.assertEqual(len(drainer._workers), 0) + + def test_drainer_shutdown_with_full_queue(self): + from unittest import mock + + from google.cloud.spanner_v1._helpers import _BoundedStreamDrainer + + drainer = _BoundedStreamDrainer(queue_size=2, worker_count=2) + drainer._ensure_started() + drainer._queue.put_nowait(mock.Mock()) + drainer._queue.put_nowait(mock.Mock()) + self.assertTrue(drainer._queue.full()) + + # Calling shutdown when queue is already full must not raise queue.Full + drainer.shutdown() + self.assertTrue(drainer._stopped) + + def test_drainer_reset_after_fork(self): + from google.cloud.spanner_v1._helpers import _BoundedStreamDrainer + + drainer = _BoundedStreamDrainer(queue_size=4, worker_count=2) + drainer._ensure_started() + self.assertTrue(drainer._started) + self.assertEqual(len(drainer._workers), 2) + + drainer._reset_after_fork() + self.assertFalse(drainer._started) + self.assertFalse(drainer._stopped) + self.assertEqual(len(drainer._workers), 0) + self.assertEqual(drainer._queue.qsize(), 0) + + def test_drainer_garbage_collection(self): + import gc + import weakref + + from google.cloud.spanner_v1._helpers import _BoundedStreamDrainer + + drainer = _BoundedStreamDrainer(queue_size=4, worker_count=1) + ref = weakref.ref(drainer) + del drainer + gc.collect() + + self.assertIsNone(ref()) + + def test_global_stream_drainer_reset_after_fork(self): + from google.cloud.spanner_v1 import _helpers + + _helpers._GLOBAL_STREAM_DRAINER._ensure_started() + self.assertTrue(_helpers._GLOBAL_STREAM_DRAINER._started) + + _helpers._GLOBAL_STREAM_DRAINER._reset_after_fork() + self.assertFalse(_helpers._GLOBAL_STREAM_DRAINER._started) + self.assertEqual(len(_helpers._GLOBAL_STREAM_DRAINER._workers), 0) + self.assertEqual(_helpers._GLOBAL_STREAM_DRAINER._queue.qsize(), 0) + + def test_module_drain_stream(self): + from unittest import mock + + from google.cloud.spanner_v1 import _helpers + + with mock.patch.object(_helpers._GLOBAL_STREAM_DRAINER, "drain") as mock_drain: + iterator = mock.Mock() + _helpers._drain_stream(iterator) + mock_drain.assert_called_once_with(iterator) diff --git a/packages/google-cloud-spanner/tests/unit/test_snapshot.py b/packages/google-cloud-spanner/tests/unit/test_snapshot.py index a5082e5b8aa1..a4cf27dc19d9 100644 --- a/packages/google-cloud-spanner/tests/unit/test_snapshot.py +++ b/packages/google-cloud-spanner/tests/unit/test_snapshot.py @@ -159,14 +159,15 @@ def _call_fut( request_id_manager=None if not session else session._database, ) - def _make_item(self, value, resume_token=b"", metadata=None): + def _make_item(self, value, resume_token=b"", metadata=None, last=False): return mock.Mock( value=value, resume_token=resume_token, metadata=metadata, precommit_token=None, + last=last, _pb=None, - spec=["value", "resume_token", "metadata", "precommit_token"], + spec=["value", "resume_token", "metadata", "precommit_token", "last"], ) def test_iteration_w_empty_raw(self): @@ -212,6 +213,226 @@ def test_iteration_w_non_empty_raw(self): ) self.assertNoSpans() + def test_restart_on_unavailable_last(self): + item_last = self._make_item(0, last=True) + trailing_item = self._make_item(1) + + raw = _MockIterator(item_last, trailing_item) + restart = mock.Mock(return_value=raw) + request = mock.Mock(test="test", spec=["test", "resume_token"]) + database = _Database() + database.spanner_api = build_spanner_api() + session = _Session(database) + derived = _build_snapshot_derived(session) + + with mock.patch("google.cloud.spanner_v1.snapshot._drain_stream") as mock_drain: + resumable = self._call_fut(derived, restart, request, session=session) + items = list(resumable) + + mock_drain.assert_called_once_with(raw) + self.assertEqual(len(items), 1) + self.assertEqual(items[0], item_last) + + def test_restart_on_unavailable_finally_cancels_on_early_termination(self): + item = self._make_item(0, last=False) + raw = _MockIterator(item) + raw.cancel = mock.Mock() + restart = mock.Mock(return_value=raw) + request = mock.Mock(test="test", spec=["test", "resume_token"]) + database = _Database() + database.spanner_api = build_spanner_api() + session = _Session(database) + derived = _build_snapshot_derived(session) + + resumable = self._call_fut(derived, restart, request, session=session) + for received_item in resumable: + break + resumable.close() + + raw.cancel.assert_called_once() + + def test_restart_on_unavailable_item_without_last_attribute(self): + item = mock.Mock( + value=0, + resume_token=b"", + metadata=None, + precommit_token=None, + _pb=None, + spec=["value", "resume_token", "metadata", "precommit_token"], + ) + raw = _MockIterator(item) + restart = mock.Mock(return_value=raw) + request = mock.Mock(test="test", spec=["test", "resume_token"]) + database = _Database() + database.spanner_api = build_spanner_api() + session = _Session(database) + derived = _build_snapshot_derived(session) + + resumable = self._call_fut(derived, restart, request, session=session) + items = list(resumable) + self.assertEqual(items, [item]) + + def test_restart_on_unavailable_finally_handles_cancel_exception(self): + item = self._make_item(0, last=False) + raw = _MockIterator(item) + raw.cancel = mock.Mock(side_effect=RuntimeError("cancel failed")) + restart = mock.Mock(return_value=raw) + request = mock.Mock(test="test", spec=["test", "resume_token"]) + database = _Database() + database.spanner_api = build_spanner_api() + session = _Session(database) + derived = _build_snapshot_derived(session) + + resumable = self._call_fut(derived, restart, request, session=session) + for _ in resumable: + break + resumable.close() + raw.cancel.assert_called_once() + + def test_restart_on_unavailable_last_does_not_cancel_iterator_in_finally(self): + item_last = self._make_item(0, last=True) + raw = _MockIterator(item_last) + raw.cancel = mock.Mock() + restart = mock.Mock(return_value=raw) + request = mock.Mock(test="test", spec=["test", "resume_token"]) + database = _Database() + database.spanner_api = build_spanner_api() + session = _Session(database) + derived = _build_snapshot_derived(session) + + with mock.patch("google.cloud.spanner_v1.snapshot._drain_stream"): + resumable = self._call_fut(derived, restart, request, session=session) + list(resumable) + + raw.cancel.assert_not_called() + + def test_streamed_result_set_with_last(self): + from google.protobuf.struct_pb2 import Value + + from google.cloud.spanner_v1.streamed import StreamedResultSet + from google.cloud.spanner_v1.types.result_set import ( + PartialResultSet, + ResultSetMetadata, + ) + from google.cloud.spanner_v1.types.type import StructType, Type, TypeCode + + fields = [StructType.Field(name="greeting", type_=Type(code=TypeCode.STRING))] + metadata_pb = ResultSetMetadata(row_type=StructType(fields=fields)) + item = PartialResultSet(metadata=metadata_pb, last=True) + item.values.append(Value(string_value="hello")) + + raw = _MockIterator(item) + streamed_result_set = StreamedResultSet(raw) + rows = list(streamed_result_set) + + self.assertEqual(rows, [["hello"]]) + self.assertTrue(streamed_result_set._done) + + def test_restart_on_unavailable_multi_chunk_with_last(self): + from google.protobuf.struct_pb2 import Value + + from google.cloud.spanner_v1.streamed import StreamedResultSet + from google.cloud.spanner_v1.types.result_set import ( + PartialResultSet, + ResultSetMetadata, + ResultSetStats, + ) + from google.cloud.spanner_v1.types.type import StructType, Type, TypeCode + + fields = [StructType.Field(name="greeting", type_=Type(code=TypeCode.STRING))] + metadata_pb = ResultSetMetadata(row_type=StructType(fields=fields)) + stats_pb = ResultSetStats(row_count_exact=2) + + chunk_one = PartialResultSet( + metadata=metadata_pb, last=False, resume_token=b"token_1" + ) + chunk_one.values.append(Value(string_value="hello")) + + chunk_two = PartialResultSet(last=True, stats=stats_pb) + chunk_two.values.append(Value(string_value="world")) + + raw = _MockIterator(chunk_one, chunk_two) + restart = mock.Mock(return_value=raw) + request = mock.Mock(test="test", spec=["test", "resume_token"]) + database = _Database() + database.spanner_api = build_spanner_api() + session = _Session(database) + derived = _build_snapshot_derived(session) + + with mock.patch("google.cloud.spanner_v1.snapshot._drain_stream") as mock_drain: + resumable = self._call_fut(derived, restart, request, session=session) + streamed_result_set = StreamedResultSet(resumable) + rows = list(streamed_result_set) + + mock_drain.assert_called_once_with(raw) + self.assertEqual(rows, [["hello"], ["world"]]) + self.assertEqual(streamed_result_set.metadata, metadata_pb) + self.assertEqual(streamed_result_set.stats, stats_pb) + self.assertTrue(streamed_result_set._done) + + def test_restart_on_unavailable_zero_rows_with_last(self): + from google.cloud.spanner_v1.streamed import StreamedResultSet + from google.cloud.spanner_v1.types.result_set import ( + PartialResultSet, + ResultSetMetadata, + ResultSetStats, + ) + from google.cloud.spanner_v1.types.type import StructType, Type, TypeCode + + fields = [StructType.Field(name="greeting", type_=Type(code=TypeCode.STRING))] + metadata_pb = ResultSetMetadata(row_type=StructType(fields=fields)) + stats_pb = ResultSetStats(row_count_exact=0) + + chunk = PartialResultSet(metadata=metadata_pb, last=True, stats=stats_pb) + + raw = _MockIterator(chunk) + restart = mock.Mock(return_value=raw) + request = mock.Mock(test="test", spec=["test", "resume_token"]) + database = _Database() + database.spanner_api = build_spanner_api() + session = _Session(database) + derived = _build_snapshot_derived(session) + + with mock.patch("google.cloud.spanner_v1.snapshot._drain_stream") as mock_drain: + resumable = self._call_fut(derived, restart, request, session=session) + streamed_result_set = StreamedResultSet(resumable) + rows = list(streamed_result_set) + + mock_drain.assert_called_once_with(raw) + self.assertEqual(rows, []) + self.assertEqual(streamed_result_set.metadata, metadata_pb) + self.assertEqual(streamed_result_set.stats, stats_pb) + self.assertTrue(streamed_result_set._done) + + def test_restart_on_unavailable_retry_before_last(self): + from google.api_core.exceptions import ServiceUnavailable + + chunk_one = self._make_item(0, resume_token=RESUME_TOKEN, last=False) + chunk_two = self._make_item(1, last=True) + + stream_one = _MockIterator( + chunk_one, fail_after=True, error=ServiceUnavailable("transient") + ) + stream_two = _MockIterator(chunk_two) + stream_two.cancel = mock.Mock() + + restart = mock.Mock(side_effect=[stream_one, stream_two]) + request = mock.Mock(test="test", spec=["test", "resume_token"]) + database = _Database() + database.spanner_api = build_spanner_api() + session = _Session(database) + derived = _build_snapshot_derived(session) + + with mock.patch("google.cloud.spanner_v1.snapshot._drain_stream") as mock_drain: + resumable = self._call_fut(derived, restart, request, session=session) + items = list(resumable) + + self.assertEqual(items, [chunk_one, chunk_two]) + self.assertEqual(len(restart.mock_calls), 2) + self.assertEqual(request.resume_token, RESUME_TOKEN) + mock_drain.assert_called_once_with(stream_two) + stream_two.cancel.assert_not_called() + def test_iteration_w_raw_w_resume_token(self): ITEMS = ( self._make_item(0),