fix: store prefetched exceptions

google/api_core/grpc_helpers.py

@@ -83,20 +83,26 @@ def error_remapped_callable(*args, **kwargs):
 class _StreamingResponseIterator(Generic[P], grpc.Call):
     def __init__(self, wrapped, prefetch_first_result=True):
         self._wrapped = wrapped
+        self._has_stored_first_result = False
 
         # This iterator is used in a retry context, and returned outside after init.
         # gRPC will not throw an exception until the stream is consumed, so we need
         # to retrieve the first result, in order to fail, in order to trigger a retry.
         try:
             if prefetch_first_result:
                 self._stored_first_result = next(self._wrapped)
+                self._has_stored_first_result = True
         except TypeError:
             # It is possible the wrapped method isn't an iterable (a grpc.Call
             # for instance). If this happens don't store the first result.
             pass
         except StopIteration:
             # ignore stop iteration at this time. This should be handled outside of retry.
             pass
+        except grpc.RpcError as exc:
+            # If the pre-fetch fails, store exception to be raised on next() call.
+            self._stored_first_error = exc
+            self._has_stored_first_result = True
 
     def __iter__(self) -> Iterator[P]:
         """This iterator is also an iterable that returns itself."""
@@ -109,10 +115,17 @@ def __next__(self) -> P:
             protobuf.Message: A single response from the stream.
         """
         try:
-            if hasattr(self, "_stored_first_result"):
-                result = self._stored_first_result
-                del self._stored_first_result
-                return result
+            if self._has_stored_first_result:
+                # if we have cached results, return or raise them on first call.
+                self._has_stored_first_result = False
+                if hasattr(self, "_stored_first_result"):
+                    result = self._stored_first_result
+                    del self._stored_first_result
+                    return result
+                if hasattr(self, "_stored_first_error"):
+                    err = self._stored_first_error
+                    del self._stored_first_error
+                    raise err
             return next(self._wrapped)
         except grpc.RpcError as exc:
             # If the stream has already returned data, we cannot recover here.

tests/unit/test_grpc_helpers.py

@@ -96,11 +96,17 @@ def test_ctor_explicit(self):
         assert list(wrapped) == ["a", "b", "c"]
 
     def test_ctor_w_rpc_error_on_prefetch(self):
+        """
+        prefetch errors should not be raised until the first iteration.
+        """
         wrapped = mock.MagicMock()
-        wrapped.__next__.side_effect = grpc.RpcError()
+        exc = grpc.RpcError()
+        wrapped.__next__.side_effect = exc
 
-        with pytest.raises(grpc.RpcError):
-            self._make_one(wrapped)
+        obj = self._make_one(wrapped)
+        assert obj._stored_first_error is exc
+        with pytest.raises(exceptions.GoogleAPICallError):
+            next(obj)
 
     def test___iter__(self):
         wrapped = self._make_wrapped("a", "b", "c")
@@ -323,11 +329,13 @@ def test_wrap_stream_errors_iterator_initialization():
 
     wrapped_callable = grpc_helpers._wrap_stream_errors(callable_)
 
-    with pytest.raises(exceptions.ServiceUnavailable) as exc_info:
-        wrapped_callable(1, 2, three="four")
+    it = wrapped_callable(1, 2, three="four")
+    assert it._stored_first_error is grpc_error
 
     callable_.assert_called_once_with(1, 2, three="four")
-    assert exc_info.value.response == grpc_error
+    # should raise on first iteration
+    with pytest.raises(exceptions.ServiceUnavailable):
+        next(it)
 
 
 def test_wrap_stream_errors_during_iteration():