From 874a8945ef397d1acb13c7077e652e3c9f8b581e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Knut=20Olav=20L=C3=B8ite?= Date: Wed, 9 Sep 2026 15:32:18 +0200 Subject: [PATCH] perf(spanner): support PartialResultSet.last with background stream draining When Cloud Spanner finishes transmitting query results, it marks `last = True` on the final PartialResultSet chunk. Previously, the client blocked synchronously waiting for gRPC trailers and EOF frames over the wire before returning the final rows. Additionally, abandoning streams early caused gRPC's C-core finalizer to mark them as CANCELLED upon garbage collection. This change enables immediate return upon observing `last = True` and offloads trailing metadata consumption to the background so streams complete cleanly with status OK: - In sync mode, completed streams are handed off to `_BoundedStreamDrainer`, which uses a bounded queue and daemon worker threads to drain to EOF. If the queue is full or the interpreter is shutting down, it falls back to inline draining. Process fork safety is ensured via `os.register_at_fork`. - In async mode, trailing frames are drained via a background `asyncio.create_task` with strong reference retention to prevent premature task garbage collection and clean cancellation handling. - Transaction precommit tokens, query stats, and metadata present on the final chunk are captured before handing the stream off to background draining. --- .../cloud/spanner_v1/_async/_helpers.py | 42 +++ .../cloud/spanner_v1/_async/snapshot.py | 189 +++++----- .../google/cloud/spanner_v1/_helpers.py | 127 +++++++ .../google/cloud/spanner_v1/snapshot.py | 152 ++++---- .../tests/unit/_async/test_helpers_extra.py | 92 +++++ .../tests/unit/_async/test_snapshot.py | 328 ++++++++++++++++++ .../tests/unit/test__helpers.py | 257 ++++++++++++++ .../tests/unit/test_snapshot.py | 225 +++++++++++- 8 files changed, 1263 insertions(+), 149 deletions(-) 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),