fix: consume part of StreamingResponseIterator to support failure while under a retry context

api_core/google/api_core/grpc_helpers.py

@@ -65,6 +65,20 @@ class _StreamingResponseIterator(grpc.Call):
     def __init__(self, wrapped):
         self._wrapped = wrapped
 
+        # 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:
+            self._stored_first_result = six.next(self._wrapped)
+        except TypeError:
+            # It is possible the wrappe 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
+
+
     def __iter__(self):
         """This iterator is also an iterable that returns itself."""
         return self
@@ -76,8 +90,13 @@ def next(self):
             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
             return six.next(self._wrapped)
         except grpc.RpcError as exc:
+            # If the stream has already returned data, we cannot recover here. 
             six.raise_from(exceptions.from_grpc_error(exc), exc)
 
     # Alias needed for Python 2/3 support.

api_core/tests/unit/test_grpc_helpers.py

@@ -88,6 +88,7 @@ def test_wrap_stream_iterable_iterface():
 
     got_iterator = wrapped_callable()
 
+
     callable_.assert_called_once_with()
 
     # Check each aliased method in the grpc.Call interface
@@ -146,10 +147,8 @@ def test_wrap_stream_errors_iterator():
 
     wrapped_callable = grpc_helpers._wrap_stream_errors(callable_)
 
-    got_iterator = wrapped_callable(1, 2, three="four")
-
     with pytest.raises(exceptions.ServiceUnavailable) as exc_info:
-        next(got_iterator)
+        got_iterator = wrapped_callable(1, 2, three="four")
 
     callable_.assert_called_once_with(1, 2, three="four")
     assert exc_info.value.response == grpc_error