Use a class to wrap grpc streaming errors instead of monkey-patching (#4995)

Author: Jon Wayne Parrott

google/api_core/grpc_helpers.py

@@ -58,6 +58,55 @@ def error_remapped_callable(*args, **kwargs):
     return error_remapped_callable
 
 
+class _StreamingResponseIterator(grpc.Call):
+    def __init__(self, wrapped):
+        self._wrapped = wrapped
+
+    def __iter__(self):
+        """This iterator is also an iterable that returns itself."""
+        return self
+
+    def next(self):
+        """Get the next response from the stream.
+
+        Returns:
+            protobuf.Message: A single response from the stream.
+        """
+        try:
+            return six.next(self._wrapped)
+        except grpc.RpcError as exc:
+            six.raise_from(exceptions.from_grpc_error(exc), exc)
+
+    # Alias needed for Python 2/3 support.
+    __next__ = next
+
+    # grpc.Call & grpc.RpcContext interface
+
+    def add_callback(self, callback):
+        return self._wrapped.add_callback(callback)
+
+    def cancel(self):
+        return self._wrapped.cancel()
+
+    def code(self):
+        return self._wrapped.code()
+
+    def details(self):
+        return self._wrapped.details()
+
+    def initial_metadata(self):
+        return self._wrapped.initial_metadata()
+
+    def is_active(self):
+        return self._wrapped.is_active()
+
+    def time_remaining(self):
+        return self._wrapped.time_remaining()
+
+    def trailing_metadata(self):
+        return self._wrapped.trailing_metadata()
+
+
 def _wrap_stream_errors(callable_):
     """Wrap errors for Unary-Stream and Stream-Stream gRPC callables.
 
@@ -71,18 +120,7 @@ def _wrap_stream_errors(callable_):
     def error_remapped_callable(*args, **kwargs):
         try:
             result = callable_(*args, **kwargs)
-            # Note: we are patching the private grpc._channel._Rendezvous._next
-            # method as magic methods (__next__ in this case) can not be
-            # patched on a per-instance basis (see
-            # https://docs.python.org/3/reference/datamodel.html
-            # #special-lookup).
-            # In an ideal world, gRPC would return a *specific* interface
-            # from *StreamMultiCallables, but they return a God class that's
-            # a combination of basically every interface in gRPC making it
-            # untenable for us to implement a wrapper object using the same
-            # interface.
-            result._next = _wrap_unary_errors(result._next)
-            return result
+            return _StreamingResponseIterator(result)
         except grpc.RpcError as exc:
             six.raise_from(exceptions.from_grpc_error(exc), exc)
 

tests/unit/test_grpc_helpers.py

@@ -66,6 +66,57 @@ def test_wrap_unary_errors():
     assert exc_info.value.response == grpc_error
 
 
+def test_wrap_stream_okay():
+    expected_responses = [1, 2, 3]
+    callable_ = mock.Mock(spec=[
+        '__call__'], return_value=iter(expected_responses))
+
+    wrapped_callable = grpc_helpers._wrap_stream_errors(callable_)
+
+    got_iterator = wrapped_callable(1, 2, three='four')
+
+    responses = list(got_iterator)
+
+    callable_.assert_called_once_with(1, 2, three='four')
+    assert responses == expected_responses
+
+
+def test_wrap_stream_iterable_iterface():
+    response_iter = mock.create_autospec(grpc.Call, instance=True)
+    callable_ = mock.Mock(spec=['__call__'], return_value=response_iter)
+
+    wrapped_callable = grpc_helpers._wrap_stream_errors(callable_)
+
+    got_iterator = wrapped_callable()
+
+    callable_.assert_called_once_with()
+
+    # Check each aliased method in the grpc.Call interface
+    got_iterator.add_callback(mock.sentinel.callback)
+    response_iter.add_callback.assert_called_once_with(mock.sentinel.callback)
+
+    got_iterator.cancel()
+    response_iter.cancel.assert_called_once_with()
+
+    got_iterator.code()
+    response_iter.code.assert_called_once_with()
+
+    got_iterator.details()
+    response_iter.details.assert_called_once_with()
+
+    got_iterator.initial_metadata()
+    response_iter.initial_metadata.assert_called_once_with()
+
+    got_iterator.is_active()
+    response_iter.is_active.assert_called_once_with()
+
+    got_iterator.time_remaining()
+    response_iter.time_remaining.assert_called_once_with()
+
+    got_iterator.trailing_metadata()
+    response_iter.trailing_metadata.assert_called_once_with()
+
+
 def test_wrap_stream_errors_invocation():
     grpc_error = RpcErrorImpl(grpc.StatusCode.INVALID_ARGUMENT)
     callable_ = mock.Mock(spec=['__call__'], side_effect=grpc_error)
@@ -83,16 +134,10 @@ class RpcResponseIteratorImpl(object):
     def __init__(self, exception):
         self._exception = exception
 
-    # Note: This matches grpc._channel._Rendezvous._next which is what is
-    # patched by _wrap_stream_errors.
-    def _next(self):
+    def next(self):
         raise self._exception
 
-    def __next__(self):  # pragma: NO COVER
-        return self._next()
-
-    def next(self):  # pragma: NO COVER
-        return self._next()
+    __next__ = next
 
 
 def test_wrap_stream_errors_iterator():
@@ -107,7 +152,6 @@ def test_wrap_stream_errors_iterator():
     with pytest.raises(exceptions.ServiceUnavailable) as exc_info:
         next(got_iterator)
 
-    assert got_iterator == response_iter
     callable_.assert_called_once_with(1, 2, three='four')
     assert exc_info.value.response == grpc_error