fix: consume part of StreamingResponseIterator to support failure while under a retry context (#10206)

1 files changed, 18 insertions, 0 deletions

diff --git a/google/api_core/grpc_helpers.py b/google/api_core/grpc_helpers.py
index 4d63beb..c47b09f 100644
--- a/google/api_core/grpc_helpers.py
+++ b/google/api_core/grpc_helpers.py
@@ -65,6 +65,19 @@ 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 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
+
     def __iter__(self):
         """This iterator is also an iterable that returns itself."""
         return self
@@ -76,8 +89,13 @@ class _StreamingResponseIterator(grpc.Call):
             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.