From 1b8a9256c02b3475d66b2aca385392a8e7906297 Mon Sep 17 00:00:00 2001 From: Subham Sinha Date: Thu, 15 Oct 2026 12:36:27 +0000 Subject: [PATCH] fix(google-cloud-spanner): record streaming operation latency and attempt latency upon stream completion Wait until the streaming response completes (either via PartialResultSet.last=True or iterator exhaustion/closure) before recording operation completion and attempt completion metrics in _restart_on_unavailable and _StreamingResponseWrapper / _AsyncStreamingResponseWrapper. Fixes #18656 --- .../cloud/spanner_v1/_async/snapshot.py | 11 + .../spanner_v1/metrics/metrics_interceptor.py | 10 +- .../google/cloud/spanner_v1/snapshot.py | 10 + .../test_operation_metrics.py | 402 ++++++++++++++++++ .../tests/system/test_metrics.py | 119 ++++-- .../tests/unit/_async/test_snapshot.py | 46 +- .../tests/unit/test_metrics_interceptor.py | 25 ++ .../tests/unit/test_snapshot.py | 32 +- 8 files changed, 617 insertions(+), 38 deletions(-) create mode 100644 packages/google-cloud-spanner/tests/mockserver_tests/test_operation_metrics.py 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 116a8a667352..d4d054aaa52c 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 @@ -203,6 +203,8 @@ async def _restart_on_unavailable( if item.resume_token: resume_token = item.resume_token break + else: + stream_finished = True except ServiceUnavailable: del item_buffer[:] @@ -238,12 +240,21 @@ async def _restart_on_unavailable( iterator = None break + if stream_finished and metrics_tracer is not None: + try: + metrics_tracer.record_operation_completion() + except Exception: + pass + finally: + metrics_tracer = None + for item in item_buffer: yield item del item_buffer[:] if stream_finished: + iterator = None break finally: try: diff --git a/packages/google-cloud-spanner/google/cloud/spanner_v1/metrics/metrics_interceptor.py b/packages/google-cloud-spanner/google/cloud/spanner_v1/metrics/metrics_interceptor.py index 8d07aae3608b..ff82f3fa46dc 100644 --- a/packages/google-cloud-spanner/google/cloud/spanner_v1/metrics/metrics_interceptor.py +++ b/packages/google-cloud-spanner/google/cloud/spanner_v1/metrics/metrics_interceptor.py @@ -272,13 +272,16 @@ def __next__(self): if self._iterator is None: self._iterator = iter(self._response) try: - return next(self._iterator) + item = next(self._iterator) except StopIteration: self._record_metrics() raise except Exception: self._record_metrics() raise + if getattr(item, "last", False) is True: + self._record_metrics() + return item def _record_metrics(self): with self._lock: @@ -495,10 +498,13 @@ async def __anext__(self): else: self._iterator = self._response try: - return await self._iterator.__anext__() + item = await self._iterator.__anext__() except StopAsyncIteration: await self._record_metrics() raise except Exception: await self._record_metrics() raise + if getattr(item, "last", False) is True: + await self._record_metrics() + return item 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 838e7589bbc4..7a5e38cdaf1c 100644 --- a/packages/google-cloud-spanner/google/cloud/spanner_v1/snapshot.py +++ b/packages/google-cloud-spanner/google/cloud/spanner_v1/snapshot.py @@ -180,6 +180,8 @@ def _restart_on_unavailable( if item.resume_token: resume_token = item.resume_token break + else: + stream_finished = True except ServiceUnavailable: del item_buffer[:] request.resume_token = resume_token @@ -211,10 +213,18 @@ def _restart_on_unavailable( if len(item_buffer) == 0: iterator = None break + if stream_finished and metrics_tracer is not None: + try: + metrics_tracer.record_operation_completion() + except Exception: + pass + finally: + metrics_tracer = None for item in item_buffer: yield item del item_buffer[:] if stream_finished: + iterator = None break finally: try: diff --git a/packages/google-cloud-spanner/tests/mockserver_tests/test_operation_metrics.py b/packages/google-cloud-spanner/tests/mockserver_tests/test_operation_metrics.py new file mode 100644 index 000000000000..11b27512e21a --- /dev/null +++ b/packages/google-cloud-spanner/tests/mockserver_tests/test_operation_metrics.py @@ -0,0 +1,402 @@ +# Copyright 2025 Google LLC All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os +import time +from unittest import mock + +from google.api_core.client_options import ClientOptions +from google.auth.credentials import AnonymousCredentials +from opentelemetry.sdk.metrics import MeterProvider +from opentelemetry.sdk.metrics.export import InMemoryMetricReader + +import google.cloud.spanner_v1.client as client_mod +from google.cloud.spanner_v1 import ( + Client, + FixedSizePool, + KeySet, + TypeCode, +) +from google.cloud.spanner_v1.metrics.metrics_interceptor import MetricsInterceptor +from google.cloud.spanner_v1.metrics.spanner_metrics_tracer_factory import ( + SpannerMetricsTracerFactory, +) +from google.cloud.spanner_v1.testing.mock_spanner import SpannerServicer +from tests.mockserver_tests.mock_server_test_base import ( + MockServerTestBase, + _make_partial_result_sets, + add_error, + add_execute_streaming_sql_results, + add_select1_result, + unavailable_status, +) + + +class TestOperationMetricsIntegration(MockServerTestBase): + def setUp(self): + super().setUp() + os.environ["SPANNER_DISABLE_BUILTIN_METRICS"] = "false" + SpannerMetricsTracerFactory._metrics_tracer_factory = None + client_mod._metrics_monitor_initialized = False + + def tearDown(self): + super().tearDown() + os.environ["SPANNER_DISABLE_BUILTIN_METRICS"] = "true" + SpannerMetricsTracerFactory._metrics_tracer_factory = None + client_mod._metrics_monitor_initialized = False + + def _setup_client_and_database(self): + client = Client( + project="p", + credentials=AnonymousCredentials(), + client_options=ClientOptions( + api_endpoint="localhost:" + str(MockServerTestBase.port), + ), + ) + instance = client.instance("test-instance") + database = instance.database( + "test-database", + pool=FixedSizePool(size=10), + enable_interceptors_in_tests=True, + ) + database._interceptors.append(MetricsInterceptor()) + database._spanner_api = None + # Pre-initialize pool sessions so they don't produce metric noise during queries + session = database._pool.get() + database._pool.put(session) + return client, database + + def _extract_metrics(self, reader, method=None): + metric_data = reader.get_metrics_data() + if not metric_data: + return {} + metrics = {} + for rm in metric_data.resource_metrics: + for sm in rm.scope_metrics: + for metric in sm.metrics: + points = list(metric.data.data_points) + if method: + points = [ + p for p in points if p.attributes.get("method") == method + ] + if points: + metrics[metric.name] = points + return metrics + + def test_operation_latency_not_recorded_before_stream_consumed(self): + fields = [("col1", TypeCode.INT64)] + partials = _make_partial_result_sets( + fields, + [ + {"values": ["1"], "last": True}, + ], + ) + add_execute_streaming_sql_results("select 1", partials) + reader = InMemoryMetricReader() + meter_provider = MeterProvider(metric_readers=[reader]) + + with ( + mock.patch( + "google.cloud.spanner_v1.metrics.metrics_tracer_factory.get_meter_provider", + return_value=meter_provider, + ), + mock.patch( + "google.cloud.spanner_v1.client.MeterProvider", + return_value=meter_provider, + ), + mock.patch( + "google.cloud.spanner_v1.client._get_spanner_emulator_host", + return_value=None, + ), + ): + _, database = self._setup_client_and_database() + with database.snapshot() as snapshot: + results = snapshot.execute_sql("select 1") + + # Verify that operational metrics for ExecuteStreamingSql have NOT been recorded yet + metrics_before = self._extract_metrics( + reader, method="Spanner.ExecuteStreamingSql" + ) + self.assertNotIn("operation_latencies", metrics_before) + self.assertNotIn("operation_count", metrics_before) + + # Consume the results while keeping `results` alive in local scope + rows = list(results) + self.assertEqual(len(rows), 1) + + # Verify that operational metrics are recorded immediately upon stream completion + # even while `results` is still alive in scope (regression test for last=True). + metrics_after = self._extract_metrics( + reader, method="Spanner.ExecuteStreamingSql" + ) + self.assertIn( + "operation_latencies", + metrics_after, + f"Metrics: {list(metrics_after.keys())}", + ) + points = metrics_after["operation_latencies"] + self.assertEqual(len(points), 1) + self.assertEqual(points[0].count, 1) + + self.assertIn("operation_count", metrics_after) + cnt_points = metrics_after["operation_count"] + self.assertEqual(len(cnt_points), 1) + self.assertEqual(cnt_points[0].value, 1) + + def test_operation_latency_measures_full_stream_duration(self): + fields = [("col1", TypeCode.INT64)] + results_cfg = [ + {"values": ["1"], "last": False}, + {"values": ["2"], "last": False}, + {"values": ["3"], "last": False}, + ] + partials = _make_partial_result_sets(fields, results_cfg) + add_execute_streaming_sql_results("select 1", partials) + + orig_get = MockServerTestBase.spanner_service.mock_spanner.get_execute_streaming_sql_results + + def delayed_get_results(sql, started_transaction): + for msg in orig_get(sql, started_transaction): + time.sleep( + 0.03 + ) # 30ms server-side streaming delay per chunk (~90ms total) + yield msg + + reader = InMemoryMetricReader() + meter_provider = MeterProvider(metric_readers=[reader]) + + with ( + mock.patch.object( + MockServerTestBase.spanner_service.mock_spanner, + "get_execute_streaming_sql_results", + side_effect=delayed_get_results, + ), + mock.patch( + "google.cloud.spanner_v1.metrics.metrics_tracer_factory.get_meter_provider", + return_value=meter_provider, + ), + mock.patch( + "google.cloud.spanner_v1.client.MeterProvider", + return_value=meter_provider, + ), + mock.patch( + "google.cloud.spanner_v1.client._get_spanner_emulator_host", + return_value=None, + ), + ): + _, database = self._setup_client_and_database() + with database.snapshot() as snapshot: + results = snapshot.execute_sql("select 1") + for _ in results: + time.sleep( + 0.06 + ) # 60ms consumer processing delay per row (180ms total) + + metrics = self._extract_metrics( + reader, method="Spanner.ExecuteStreamingSql" + ) + self.assertIn("operation_latencies", metrics) + point = metrics["operation_latencies"][0] + # Operational latency must include the ~90ms server stream duration (>= 75ms) + # and exclude the 180ms consumer sleep (<= 165ms). + self.assertGreaterEqual(point.sum, 75) + self.assertLessEqual(point.sum, 165) + + def test_operation_latency_last_true_and_last_false_consistency(self): + """Verifies last=True and last=False produce consistent operation_latencies + and do not diverge when consumer sleeps while iterating rows.""" + fields = [("col1", TypeCode.INT64)] + partials_no_last = _make_partial_result_sets( + fields, + [ + {"values": ["1"], "last": False}, + {"values": ["2"], "last": False}, + {"values": ["3"], "last": False}, + ], + ) + partials_with_last = _make_partial_result_sets( + fields, + [ + {"values": ["1"], "last": False}, + {"values": ["2"], "last": False}, + {"values": ["3"], "last": True}, + ], + ) + add_execute_streaming_sql_results("select no_last", partials_no_last) + add_execute_streaming_sql_results("select with_last", partials_with_last) + + orig_get = MockServerTestBase.spanner_service.mock_spanner.get_execute_streaming_sql_results + + def delayed_get_results(sql, started_transaction): + for msg in orig_get(sql, started_transaction): + time.sleep(0.025) # 25ms server streaming delay per chunk (~75ms total) + yield msg + + def run_query_and_get_latency(sql): + SpannerMetricsTracerFactory._metrics_tracer_factory = None + client_mod._metrics_monitor_initialized = False + reader = InMemoryMetricReader() + meter_provider = MeterProvider(metric_readers=[reader]) + with ( + mock.patch.object( + MockServerTestBase.spanner_service.mock_spanner, + "get_execute_streaming_sql_results", + side_effect=delayed_get_results, + ), + mock.patch( + "google.cloud.spanner_v1.metrics.metrics_tracer_factory.get_meter_provider", + return_value=meter_provider, + ), + mock.patch( + "google.cloud.spanner_v1.client.MeterProvider", + return_value=meter_provider, + ), + mock.patch( + "google.cloud.spanner_v1.client._get_spanner_emulator_host", + return_value=None, + ), + ): + _, database = self._setup_client_and_database() + with database.snapshot() as snapshot: + results = snapshot.execute_sql(sql) + for _ in results: + time.sleep(0.06) # 60ms/row consumer sleep (180ms total) + metrics = self._extract_metrics( + reader, method="Spanner.ExecuteStreamingSql" + ) + point = metrics["operation_latencies"][0] + self.assertEqual(point.count, 1) + return point.sum + + lat_no_last = run_query_and_get_latency("select no_last") + lat_with_last = run_query_and_get_latency("select with_last") + + # Both must reflect the ~75ms server streaming time and exclude the 180ms consumer sleep + self.assertGreaterEqual(lat_no_last, 60) + self.assertLessEqual(lat_no_last, 140) + self.assertGreaterEqual(lat_with_last, 60) + self.assertLessEqual(lat_with_last, 140) + + def test_operation_latency_with_retries(self): + add_select1_result() + add_error(SpannerServicer.ExecuteStreamingSql.__name__, unavailable_status()) + reader = InMemoryMetricReader() + meter_provider = MeterProvider(metric_readers=[reader]) + + with ( + mock.patch( + "google.cloud.spanner_v1.metrics.metrics_tracer_factory.get_meter_provider", + return_value=meter_provider, + ), + mock.patch( + "google.cloud.spanner_v1.client.MeterProvider", + return_value=meter_provider, + ), + mock.patch( + "google.cloud.spanner_v1.client._get_spanner_emulator_host", + return_value=None, + ), + ): + _, database = self._setup_client_and_database() + with database.snapshot() as snapshot: + results = snapshot.execute_sql("select 1") + rows = list(results) + self.assertEqual(len(rows), 1) + + metrics = self._extract_metrics( + reader, method="Spanner.ExecuteStreamingSql" + ) + # Should record 1 logical operation but 2 attempts + self.assertIn("operation_count", metrics) + op_cnt_point = metrics["operation_count"][0] + self.assertEqual(op_cnt_point.value, 1) + + self.assertIn("attempt_count", metrics) + att_cnt_point = metrics["attempt_count"][0] + self.assertEqual(att_cnt_point.value, 2) + + self.assertIn("operation_latencies", metrics) + op_lat_point = metrics["operation_latencies"][0] + self.assertEqual(op_lat_point.count, 1) + + def test_session_ping_emits_operation_metrics(self): + add_select1_result() + reader = InMemoryMetricReader() + meter_provider = MeterProvider(metric_readers=[reader]) + + with ( + mock.patch( + "google.cloud.spanner_v1.metrics.metrics_tracer_factory.get_meter_provider", + return_value=meter_provider, + ), + mock.patch( + "google.cloud.spanner_v1.client.MeterProvider", + return_value=meter_provider, + ), + mock.patch( + "google.cloud.spanner_v1.client._get_spanner_emulator_host", + return_value=None, + ), + ): + _, database = self._setup_client_and_database() + session = database._pool.get() + try: + session.ping() + finally: + database._pool.put(session) + + metrics = self._extract_metrics(reader, method="Spanner.ExecuteSql") + self.assertIn("operation_latencies", metrics) + self.assertEqual(metrics["operation_latencies"][0].count, 1) + self.assertIn("operation_count", metrics) + self.assertEqual(metrics["operation_count"][0].value, 1) + + def test_operation_latency_streaming_read(self): + reader = InMemoryMetricReader() + meter_provider = MeterProvider(metric_readers=[reader]) + + with ( + mock.patch( + "google.cloud.spanner_v1.metrics.metrics_tracer_factory.get_meter_provider", + return_value=meter_provider, + ), + mock.patch( + "google.cloud.spanner_v1.client.MeterProvider", + return_value=meter_provider, + ), + mock.patch( + "google.cloud.spanner_v1.client._get_spanner_emulator_host", + return_value=None, + ), + ): + _, database = self._setup_client_and_database() + with database.snapshot() as snapshot: + results = snapshot.read( + table="citizens", + columns=["email", "first_name"], + keyset=KeySet(all_=True), + ) + # Before consumption + metrics_before = self._extract_metrics( + reader, method="Spanner.StreamingRead" + ) + self.assertNotIn("operation_latencies", metrics_before) + + # Consume + list(results) + + metrics = self._extract_metrics(reader, method="Spanner.StreamingRead") + self.assertIn("operation_latencies", metrics) + op_lat_point = metrics["operation_latencies"][0] + self.assertEqual(op_lat_point.count, 1) diff --git a/packages/google-cloud-spanner/tests/system/test_metrics.py b/packages/google-cloud-spanner/tests/system/test_metrics.py index 71a51a733205..909ceacda9af 100644 --- a/packages/google-cloud-spanner/tests/system/test_metrics.py +++ b/packages/google-cloud-spanner/tests/system/test_metrics.py @@ -19,29 +19,12 @@ from opentelemetry.sdk.metrics import MeterProvider from opentelemetry.sdk.metrics.export import InMemoryMetricReader -from google.cloud.spanner_v1._helpers import _GLOBAL_STREAM_DRAINER from google.cloud.spanner_v1.metrics import constants from google.cloud.spanner_v1.metrics.spanner_metrics_tracer_factory import ( SpannerMetricsTracerFactory, ) -def _wait_for_stream_drainer(timeout: float = 10.0) -> None: - """Wait for pending background stream draining tasks to complete. - - TODO(#18656): Remove this helper once metrics are recorded immediately - upon observing PartialResultSet.last=True. - """ - drain_queue = _GLOBAL_STREAM_DRAINER._queue - with drain_queue.all_tasks_done: - drained = drain_queue.all_tasks_done.wait_for( - lambda: drain_queue.unfinished_tasks == 0, - timeout=timeout, - ) - if not drained: - pytest.fail("Timed out waiting for background stream drainer to complete.") - - def test_builtin_metrics_with_default_otel(not_emulator, shared_database): """ Verifies that built-in metrics are collected by default when a @@ -67,13 +50,6 @@ def test_builtin_metrics_with_default_otel(not_emulator, shared_database): ): with shared_database.snapshot() as snapshot: list(snapshot.execute_sql("SELECT 1")) - # When PartialResultSet.last is True, stream draining and metric recording - # for attempt/GFE/AFE latencies happen asynchronously in the background. - # Wait for any pending drain tasks to complete before collecting metrics. - # TODO(#18656): Record attempt and front-end metrics when PartialResultSet.last=True - # is observed on the stream, rather than waiting for StopIteration in the background - # drainer. This synchronization step can then be removed. - _wait_for_stream_drainer() finally: SpannerMetricsTracerFactory._metrics_tracer_factory = old_factory @@ -104,3 +80,98 @@ def test_builtin_metrics_with_default_otel(not_emulator, shared_database): return pytest.fail(f"Metric '{constants.METRIC_NAME_OPERATION_COUNT}' not found.") + + +def test_operation_latency_streaming_lifecycle(not_emulator, shared_database): + """ + Verifies that operation latency is not recorded prematurely before stream + consumption, completes immediately when the stream finishes (even while + the StreamedResultSet reference remains in scope), and excludes consumer sleep. + """ + import time + + reader = InMemoryMetricReader() + meter_provider = MeterProvider(metric_readers=[reader]) + + old_factory = SpannerMetricsTracerFactory._metrics_tracer_factory + SpannerMetricsTracerFactory._metrics_tracer_factory = None + try: + with ( + mock.patch( + "google.cloud.spanner_v1.metrics.metrics_tracer_factory.get_meter_provider", + return_value=meter_provider, + ), + mock.patch( + "google.cloud.spanner_v1.client.MeterProvider", + return_value=meter_provider, + ), + mock.patch.dict(os.environ, {"SPANNER_DISABLE_BUILTIN_METRICS": "false"}), + ): + with shared_database.snapshot() as snapshot: + results = snapshot.execute_sql("SELECT 1") + + # Before consumption, operational metrics for the stream should not be completed + metric_data_before = reader.get_metrics_data() + streaming_op_latencies_before = [ + p + for rm in ( + metric_data_before.resource_metrics + if metric_data_before + else [] + ) + for sm in rm.scope_metrics + for metric in sm.metrics + if metric.name + in ( + constants.METRIC_NAME_OPERATION_LATENCIES, + "operation_latencies", + ) + for p in metric.data.data_points + if p.attributes.get("method") + in ("ExecuteStreamingSql", "Spanner.ExecuteStreamingSql") + ] + assert len(streaming_op_latencies_before) == 0 + + # Consume stream; consumer sleep should not inflate operation_latencies + for _ in results: + time.sleep(0.1) + + # Verify metrics are recorded while `results` is still in scope + metric_data_after = reader.get_metrics_data() + finally: + SpannerMetricsTracerFactory._metrics_tracer_factory = old_factory + + assert metric_data_after is not None + + streaming_op_latencies_after = [ + p + for rm in metric_data_after.resource_metrics + for sm in rm.scope_metrics + for metric in sm.metrics + if metric.name + in (constants.METRIC_NAME_OPERATION_LATENCIES, "operation_latencies") + for p in metric.data.data_points + if p.attributes.get("method") + in ("ExecuteStreamingSql", "Spanner.ExecuteStreamingSql") + ] + streaming_attempt_latencies_after = [ + p + for rm in metric_data_after.resource_metrics + for sm in rm.scope_metrics + for metric in sm.metrics + if metric.name in (constants.METRIC_NAME_ATTEMPT_LATENCIES, "attempt_latencies") + for p in metric.data.data_points + if p.attributes.get("method") + in ("ExecuteStreamingSql", "Spanner.ExecuteStreamingSql") + ] + assert len(streaming_op_latencies_after) == 1 + assert streaming_op_latencies_after[0].sum > 0 + assert len(streaming_attempt_latencies_after) > 0 + assert ( + streaming_op_latencies_after[0].sum >= streaming_attempt_latencies_after[0].sum + ) + # Operation latency should be close to attempt latency, excluding the 100ms consumer sleep: + assert ( + streaming_op_latencies_after[0].sum - streaming_attempt_latencies_after[0].sum + < 50 + ) 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 f85a4c49b1b4..3fc7f3cf13aa 100644 --- a/packages/google-cloud-spanner/tests/unit/_async/test_snapshot.py +++ b/packages/google-cloud-spanner/tests/unit/_async/test_snapshot.py @@ -772,7 +772,7 @@ async def test_restart_on_unavailable_early_break_no_context_error(self): item_one = mock.Mock( value=0, - resume_token=b"", + resume_token=b"r1", metadata=None, precommit_token=None, last=False, @@ -829,8 +829,17 @@ async def test_restart_on_unavailable_real_factory_early_break_no_context_error( SpannerMetricsTracerFactory, ) - item = mock.Mock( + item_one = mock.Mock( value=0, + resume_token=b"r1", + metadata=None, + precommit_token=None, + last=False, + _pb=None, + spec=["value", "resume_token", "metadata", "precommit_token", "last"], + ) + item_two = mock.Mock( + value=1, resume_token=b"", metadata=None, precommit_token=None, @@ -838,7 +847,7 @@ async def test_restart_on_unavailable_real_factory_early_break_no_context_error( _pb=None, spec=["value", "resume_token", "metadata", "precommit_token", "last"], ) - raw = _MockIterator(item) + raw = _MockIterator(item_one, item_two) restart = mock.Mock(return_value=raw) request = mock.Mock(test="test", spec=["test", "resume_token"]) database = _Database() @@ -857,7 +866,7 @@ async def test_restart_on_unavailable_real_factory_early_break_no_context_error( # Breaking early must not raise ValueError during aclose() async for it in resumable: - self.assertEqual(it, item) + self.assertEqual(it, item_one) break await resumable.aclose() @@ -976,8 +985,17 @@ async def test_restart_on_unavailable_metrics_completion_error_handled(self): "metrics export failure" ) - item = mock.Mock( + item_one = mock.Mock( value=0, + resume_token=b"r1", + metadata=None, + precommit_token=None, + last=False, + _pb=None, + spec=["value", "resume_token", "metadata", "precommit_token", "last"], + ) + item_two = mock.Mock( + value=1, resume_token=b"", metadata=None, precommit_token=None, @@ -985,7 +1003,7 @@ async def test_restart_on_unavailable_metrics_completion_error_handled(self): _pb=None, spec=["value", "resume_token", "metadata", "precommit_token", "last"], ) - mock_iterator = _MockIterator(item) + mock_iterator = _MockIterator(item_one, item_two) mock_iterator.cancel = mock.Mock() request = mock.Mock(test="test", spec=["test", "resume_token"]) database = _Database() @@ -1009,10 +1027,24 @@ async def test_restart_on_unavailable_metrics_completion_error_handled(self): break await resumable.aclose() - self.assertEqual(items, [item]) + self.assertEqual(items, [item_one]) mock_tracer.record_operation_completion.assert_called_once() mock_iterator.cancel.assert_called_once() + # Also verify error is handled when stream finishes normally before yield + mock_tracer.reset_mock() + resumable_done = _restart_on_unavailable( + mock.Mock(return_value=_MockIterator(item_two)), + request, + session=session, + transaction=snapshot, + request_id_manager=session._database, + resource_info={"project": "p"}, + ) + items_done = [it async for it in resumable_done] + self.assertEqual(items_done, [item_two]) + mock_tracer.record_operation_completion.assert_called_once() + async def test_restart_on_unavailable_metrics_completion_error_preserves_query_error( self, ): diff --git a/packages/google-cloud-spanner/tests/unit/test_metrics_interceptor.py b/packages/google-cloud-spanner/tests/unit/test_metrics_interceptor.py index 33c5ada9911a..adc69a7e198f 100644 --- a/packages/google-cloud-spanner/tests/unit/test_metrics_interceptor.py +++ b/packages/google-cloud-spanner/tests/unit/test_metrics_interceptor.py @@ -1548,3 +1548,28 @@ def test_sync_streaming_response_wrapper_del_unrecorded(mock_tracer_ctx): wrapper = _StreamingResponseWrapper(mock_response, mock_tracer_ctx) wrapper.__del__() mock_tracer_ctx.record_attempt_completion.assert_called_once() + + +def test_streaming_response_wrapper_last_true_records_metrics(mock_tracer_ctx): + mock_item = Mock(last=True) + wrapper = _StreamingResponseWrapper(iter([mock_item]), mock_tracer_ctx) + assert next(wrapper) is mock_item + mock_tracer_ctx.record_attempt_completion.assert_called_once() + + +@pytest.mark.asyncio +async def test_async_streaming_response_wrapper_last_true_records_metrics( + mock_tracer_ctx, +): + mock_item = Mock(last=True) + + class AsyncCallWithLast: + def __aiter__(self): + async def _gen(): + yield mock_item + + return _gen() + + wrapper = _AsyncStreamingResponseWrapper(AsyncCallWithLast(), mock_tracer_ctx) + assert await wrapper.__anext__() is mock_item + mock_tracer_ctx.record_attempt_completion.assert_called_once() diff --git a/packages/google-cloud-spanner/tests/unit/test_snapshot.py b/packages/google-cloud-spanner/tests/unit/test_snapshot.py index a367b293f798..9be7312ca020 100644 --- a/packages/google-cloud-spanner/tests/unit/test_snapshot.py +++ b/packages/google-cloud-spanner/tests/unit/test_snapshot.py @@ -1026,7 +1026,7 @@ def test_restart_on_unavailable_early_break_no_context_error(self): mock_tracer = mock.Mock() - FIRST = (self._make_item(0), self._make_item(1)) + FIRST = (self._make_item(0, resume_token=b"r1"), self._make_item(1)) raw = _MockIterator(*FIRST) restart = mock.Mock(return_value=raw) request = mock.Mock(test="test", spec=["test", "resume_token"]) @@ -1062,7 +1062,7 @@ def test_restart_on_unavailable_real_factory_early_break_no_context_error(self): ) from google.cloud.spanner_v1.snapshot import _restart_on_unavailable - FIRST = (self._make_item(0), self._make_item(1)) + FIRST = (self._make_item(0, resume_token=b"r1"), self._make_item(1)) raw = _MockIterator(*FIRST) restart = mock.Mock(return_value=raw) request = mock.Mock(test="test", spec=["test", "resume_token"]) @@ -1194,8 +1194,17 @@ def test_restart_on_unavailable_metrics_completion_error_handled(self): "metrics export failure" ) - item = mock.Mock( + item_one = mock.Mock( value=0, + resume_token=b"r1", + metadata=None, + precommit_token=None, + last=False, + _pb=None, + spec=["value", "resume_token", "metadata", "precommit_token", "last"], + ) + item_two = mock.Mock( + value=1, resume_token=b"", metadata=None, precommit_token=None, @@ -1203,7 +1212,7 @@ def test_restart_on_unavailable_metrics_completion_error_handled(self): _pb=None, spec=["value", "resume_token", "metadata", "precommit_token", "last"], ) - mock_iterator = _MockIterator(item) + mock_iterator = _MockIterator(item_one, item_two) mock_iterator.cancel = mock.Mock() request = mock.Mock(test="test", spec=["test", "resume_token"]) database = _Database() @@ -1228,10 +1237,23 @@ def test_restart_on_unavailable_metrics_completion_error_handled(self): break resumable.close() - self.assertEqual(items, [item]) + self.assertEqual(items, [item_one]) mock_tracer.record_operation_completion.assert_called_once() mock_iterator.cancel.assert_called_once() + # Also verify error is handled when stream finishes normally before yield + mock_tracer.reset_mock() + resumable_done = _restart_on_unavailable( + mock.Mock(return_value=_MockIterator(item_two)), + request, + session=session, + transaction=derived, + request_id_manager=session._database, + resource_info={"project": "p"}, + ) + self.assertEqual(list(resumable_done), [item_two]) + mock_tracer.record_operation_completion.assert_called_once() + def test_restart_on_unavailable_metrics_completion_error_preserves_query_error( self, ):