fix(spanner): handle errors during stream restart in snapshot (#1471)

sinhasubham · web-flow · commit c0668735cb69 · 2026-01-14T20:33:02.000+05:30
***Handle errors during stream restart in snapshot***

**Root Cause**
When `_restart_on_unavailable` caught a `ServiceUnavailable` or
resumable `InternalServerError`, it attempted to re-initialize the
iterator immediately within the `except` block. If this
re-initialization failed (e.g. due to a persistent transient error), the
exception would propagate unhandled, breaking the retry loop.

**Fix**
This change modifies the logic to reset the iterator to `None` and
`continue` the loop, forcing the re-initialization to occur inside the
`try` block. This ensures that subsequent errors during restart are
properly caught and retried.

**Testing**
Added unit tests to cover this specific behavior
diff --git a/google/cloud/spanner_v1/snapshot.py b/google/cloud/spanner_v1/snapshot.py
@@ -146,27 +146,12 @@ def _restart_on_unavailable(
 
         except ServiceUnavailable:
             del item_buffer[:]
-            with trace_call(
-                trace_name,
-                session,
-                attributes,
-                observability_options=observability_options,
-                metadata=metadata,
-            ) as span, MetricsCapture():
-                request.resume_token = resume_token
-                if transaction is not None:
-                    transaction_selector = transaction._build_transaction_selector_pb()
-                request.transaction = transaction_selector
-                attempt += 1
-                iterator = method(
-                    request=request,
-                    metadata=request_id_manager.metadata_with_request_id(
-                        nth_request,
-                        attempt,
-                        metadata,
-                        span,
-                    ),
-                )
+            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:
@@ -177,27 +162,12 @@ def _restart_on_unavailable(
             if not resumable_error:
                 raise
             del item_buffer[:]
-            with trace_call(
-                trace_name,
-                session,
-                attributes,
-                observability_options=observability_options,
-                metadata=metadata,
-            ) as span, MetricsCapture():
-                request.resume_token = resume_token
-                if transaction is not None:
-                    transaction_selector = transaction._build_transaction_selector_pb()
-                attempt += 1
-                request.transaction = transaction_selector
-                iterator = method(
-                    request=request,
-                    metadata=request_id_manager.metadata_with_request_id(
-                        nth_request,
-                        attempt,
-                        metadata,
-                        span,
-                    ),
-                )
+            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
 
         if len(item_buffer) == 0:
diff --git a/tests/unit/test_snapshot.py b/tests/unit/test_snapshot.py
@@ -405,6 +405,56 @@ def test_iteration_w_raw_raising_unavailable_after_token(self):
         self.assertEqual(request.resume_token, RESUME_TOKEN)
         self.assertNoSpans()
 
+    def test_iteration_w_raw_raising_unavailable_during_restart(self):
+        from google.api_core.exceptions import ServiceUnavailable
+
+        FIRST = (self._make_item(0), self._make_item(1, resume_token=RESUME_TOKEN))
+        LAST = (self._make_item(2),)
+        before = _MockIterator(
+            *FIRST, fail_after=True, error=ServiceUnavailable("testing")
+        )
+        after = _MockIterator(*LAST)
+        request = mock.Mock(test="test", spec=["test", "resume_token"])
+        # The second call (the first retry) raises ServiceUnavailable immediately.
+        # The third call (the second retry) succeeds.
+        restart = mock.Mock(
+            spec=[],
+            side_effect=[before, ServiceUnavailable("retry failed"), after],
+        )
+        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)
+        self.assertEqual(list(resumable), list(FIRST + LAST))
+        self.assertEqual(len(restart.mock_calls), 3)
+        self.assertEqual(request.resume_token, RESUME_TOKEN)
+        self.assertNoSpans()
+
+    def test_iteration_w_raw_raising_resumable_internal_error_during_restart(self):
+        FIRST = (self._make_item(0), self._make_item(1, resume_token=RESUME_TOKEN))
+        LAST = (self._make_item(2),)
+        before = _MockIterator(
+            *FIRST,
+            fail_after=True,
+            error=INTERNAL_SERVER_ERROR_UNEXPECTED_EOS,
+        )
+        after = _MockIterator(*LAST)
+        request = mock.Mock(test="test", spec=["test", "resume_token"])
+        restart = mock.Mock(
+            spec=[],
+            side_effect=[before, INTERNAL_SERVER_ERROR_UNEXPECTED_EOS, after],
+        )
+        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)
+        self.assertEqual(list(resumable), list(FIRST + LAST))
+        self.assertEqual(len(restart.mock_calls), 3)
+        self.assertEqual(request.resume_token, RESUME_TOKEN)
+        self.assertNoSpans()
+
     def test_iteration_w_raw_w_multiuse(self):
         from google.cloud.spanner_v1 import (
             ReadRequest,
