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

***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

Author: sinhasubham

google/cloud/spanner_v1/snapshot.py

@@ -146,27 +146,12 @@ def _restart_on_unavailable(
 
 exceptServiceUnavailable:
 delitem_buffer[:]
-withtrace_call(
-trace_name,
-session,
-attributes,
-observability_options=observability_options,
-metadata=metadata,
- ) asspan, MetricsCapture():
-request.resume_token=resume_token
-iftransactionisnotNone:
-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
+iftransactionisnotNone:
+transaction_selector=transaction._build_transaction_selector_pb()
+request.transaction=transaction_selector
+attempt+=1
+iterator=None
 continue
 
 exceptInternalServerErrorasexc:
@@ -177,27 +162,12 @@ def _restart_on_unavailable(
 ifnotresumable_error:
 raise
 delitem_buffer[:]
-withtrace_call(
-trace_name,
-session,
-attributes,
-observability_options=observability_options,
-metadata=metadata,
- ) asspan, MetricsCapture():
-request.resume_token=resume_token
-iftransactionisnotNone:
-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
+iftransactionisnotNone:
+transaction_selector=transaction._build_transaction_selector_pb()
+attempt+=1
+request.transaction=transaction_selector
+iterator=None
 continue
 
 iflen(item_buffer) ==0:

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()
 
+deftest_iteration_w_raw_raising_unavailable_during_restart(self):
+fromgoogle.api_core.exceptionsimportServiceUnavailable
+
+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()
+
+deftest_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()
+
 deftest_iteration_w_raw_w_multiuse(self):
 fromgoogle.cloud.spanner_v1import (
 ReadRequest,