Add bidiBase, refactor bidiRpc and add AsyncBidi

Author: chandra-siri

google/api_core/bidi.py

@@ -22,6 +22,7 @@
 import time
 
 from google.api_core import exceptions
+from google.api_core.bidi_base import BidiRpcBase
 
 _LOGGER = logging.getLogger(__name__)
 _BIDIRECTIONAL_CONSUMER_NAME = "Thread-ConsumeBidirectionalStream"
@@ -201,7 +202,7 @@ def __repr__(self):
         )
 
 
-class BidiRpc(object):
+class BidiRpc(BidiRpcBase):
     """A helper for consuming a bi-directional streaming RPC.
 
     This maps gRPC's built-in interface which uses a request iterator and a
@@ -240,35 +241,9 @@ class BidiRpc(object):
             the request.
     """
 
-    def __init__(self, start_rpc, initial_request=None, metadata=None):
-        self._start_rpc = start_rpc
-        self._initial_request = initial_request
-        self._rpc_metadata = metadata
-        self._request_queue = queue_module.Queue()
-        self._request_generator = None
-        self._is_active = False
-        self._callbacks = []
-        self.call = None
-
-    def add_done_callback(self, callback):
-        """Adds a callback that will be called when the RPC terminates.
-
-        This occurs when the RPC errors or is successfully terminated.
-
-        Args:
-            callback (Callable[[grpc.Future], None]): The callback to execute.
-                It will be provided with the same gRPC future as the underlying
-                stream which will also be a :class:`grpc.Call`.
-        """
-        self._callbacks.append(callback)
-
-    def _on_call_done(self, future):
-        # This occurs when the RPC errors or is successfully terminated.
-        # Note that grpc's "future" here can also be a grpc.RpcError.
-        # See note in https://github.com/grpc/grpc/issues/10885#issuecomment-302651331
-        # that `grpc.RpcError` is also `grpc.call`.
-        for callback in self._callbacks:
-            callback(future)
+    def _create_queue(self):
+        """Create a queue for requests."""
+        return queue_module.Queue()
 
     def open(self):
         """Opens the stream."""
@@ -352,11 +327,6 @@ def is_active(self):
         """bool: True if this stream is currently open and active."""
         return self.call is not None and self.call.is_active()
 
-    @property
-    def pending_requests(self):
-        """int: Returns an estimate of the number of queued requests."""
-        return self._request_queue.qsize()
-
 
 def _never_terminate(future_or_error):
     """By default, no errors cause BiDi termination."""

google/api_core/bidi_async.py

@@ -1,4 +1,4 @@
-# Copyright 2020, Google LLC
+# Copyright 2024, Google LLC
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -18,6 +18,7 @@
 import logging
 
 from google.api_core import exceptions
+from google.api_core.bidi_base import BidiRpcBase
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -69,7 +70,7 @@ class _AsyncRequestQueueGenerator:
                 request.
     """
 
-    def __init__(self, queue: asyncio.Queue, period: float = 1, initial_request=None):
+    def __init__(self, queue: asyncio.Queue, period: float = 0.1, initial_request=None):
         self._queue = queue
         self._period = period
         self._initial_request = initial_request
@@ -122,7 +123,7 @@ async def __aiter__(self):
             yield item
 
 
-class AsyncBidiRpc:
+class AsyncBidiRpc(BidiRpcBase):
     """A helper for consuming a async bi-directional streaming RPC.
 
     This maps gRPC's built-in interface which uses a request iterator and a
@@ -161,35 +162,9 @@ class AsyncBidiRpc:
             the request.
     """
 
-    def __init__(self, start_rpc, initial_request=None, metadata=None):
-        self._start_rpc = start_rpc
-        self._initial_request = initial_request
-        self._rpc_metadata = metadata
-        self._request_queue = asyncio.Queue()
-        self._request_generator = None
-        self._callbacks = []
-        self.call = None
-        self._loop = asyncio.get_event_loop()
-
-    def add_done_callback(self, callback):
-        """Adds a callback that will be called when the RPC terminates.
-
-        This occurs when the RPC errors or is successfully terminated.
-
-        Args:
-            callback (Callable[[grpc.Future], None]): The callback to execute.
-                It will be provided with the same gRPC future as the underlying
-                stream which will also be a :class:`grpc.aio.Call`.
-        """
-        self._callbacks.append(callback)
-
-    def _on_call_done(self, future):
-        # This occurs when the RPC errors or is successfully terminated.
-        # Note that grpc's "future" here can also be a grpc.RpcError.
-        # See note in https://github.com/grpc/grpc/issues/10885#issuecomment-302651331
-        # that `grpc.RpcError` is also `grpc.aio.Call`.
-        for callback in self._callbacks:
-            callback(future)
+    def _create_queue(self):
+        """Create a queue for requests."""
+        return asyncio.Queue()
 
     async def open(self):
         """Opens the stream."""
@@ -268,8 +243,3 @@ async def recv(self):
     def is_active(self):
         """bool: True if this stream is currently open and active."""
         return self.call is not None and not self.call.done()
-
-    @property
-    def pending_requests(self):
-        """int: Returns an estimate of the number of queued requests."""
-        return self._request_queue.qsize()

google/api_core/bidi_base.py

@@ -0,0 +1,95 @@
+# Copyright 2024, Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# You may obtain a copy of the License at
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Base class for Bi-directional streaming RPC helpers."""
+
+
+class BidiRpcBase:
+    """A base helper for consuming a bi-directional streaming RPC.
+
+    This maps gRPC's built-in interface which uses a request iterator and a
+    response iterator into a socket-like :func:`send` and :func:`recv`. This
+    is a more useful pattern for long-running or asymmetric streams (streams
+    where there is not a direct correlation between the requests and
+    responses).
+
+    This does *not* retry the stream on errors.
+
+    Args:
+        start_rpc (grpc.StreamStreamMultiCallable): The gRPC method used to
+            start the RPC.
+        initial_request (Union[protobuf.Message,
+                Callable[[], protobuf.Message]]): The initial request to
+            yield. This is useful if an initial request is needed to start the
+            stream.
+        metadata (Sequence[Tuple(str, str)]): RPC metadata to include in
+            the request.
+    """
+
+    def __init__(self, start_rpc, initial_request=None, metadata=None):
+        self._start_rpc = start_rpc
+        self._initial_request = initial_request
+        self._rpc_metadata = metadata
+        self._request_queue = self._create_queue()
+        self._request_generator = None
+        self._callbacks = []
+        self.call = None
+
+    def _create_queue(self):
+        """Create a queue for requests."""
+        raise NotImplementedError("Not implemented in base class")
+
+    def add_done_callback(self, callback):
+        """Adds a callback that will be called when the RPC terminates.
+
+        This occurs when the RPC errors or is successfully terminated.
+
+        Args:
+            callback (Callable[[grpc.Future], None]): The callback to execute.
+                It will be provided with the same gRPC future as the underlying
+                stream which will also be a :class:`grpc.aio.Call`.
+        """
+        self._callbacks.append(callback)
+
+    def _on_call_done(self, future):
+        # This occurs when the RPC errors or is successfully terminated.
+        # Note that grpc's "future" here can also be a grpc.RpcError.
+        # See note in https://github.com/grpc/grpc/issues/10885#issuecomment-302651331
+        # that `grpc.RpcError` is also `grpc.aio.Call`.
+        for callback in self._callbacks:
+            callback(future)
+
+    def open(self):
+        """Opens the stream."""
+        raise NotImplementedError("Not implemented in base class")
+
+    def close(self):
+        """Closes the stream."""
+        raise NotImplementedError("Not implemented in base class")
+
+    def send(self, request):
+        """Queue a message to be sent on the stream."""
+        raise NotImplementedError("Not implemented in base class")
+
+    def recv(self):
+        """Wait for a message to be returned from the stream."""
+        raise NotImplementedError("Not implemented in base class")
+
+    @property
+    def is_active(self):
+        """bool: True if this stream is currently open and active."""
+        raise NotImplementedError("Not implemented in base class")
+
+    @property
+    def pending_requests(self):
+        """int: Returns an estimate of the number of queued requests."""
+        return self._request_queue.qsize()

tests/asyncio/test_bidi_async.py

@@ -14,10 +14,12 @@
 
 import asyncio
 
+from unittest import mock
+
 try:
-    from unittest import mock
+    from unittest.mock import AsyncMock
 except ImportError:  # pragma: NO COVER
-    import mock  # type: ignore
+    from asyncmock import AsyncMock
 
 import pytest
 
@@ -138,9 +140,9 @@ async def test_exit_with_stop(self):
 def make_async_rpc():
     """Makes a mock async RPC used to test Bidi classes."""
     call = mock.create_autospec(aio.StreamStreamCall, instance=True)
-    rpc = mock.AsyncMock(spec=aio.StreamStreamMultiCallable)
+    rpc = AsyncMock()
 
-    async def rpc_side_effect(request, metadata=None):
+    def rpc_side_effect(request, metadata=None):
         call.done.return_value = False
         return call
 
@@ -151,7 +153,7 @@ def cancel_side_effect():
         return True
 
     call.cancel.side_effect = cancel_side_effect
-    call.read = mock.AsyncMock()
+    call.read = AsyncMock()
 
     return rpc, call
 
@@ -167,7 +169,6 @@ async def read(self):
         raise self.exception
 
 
-@pytest.mark.asyncio
 class TestAsyncBidiRpc:
     def test_initial_state(self):
         bidi_rpc = bidi_async.AsyncBidiRpc(None)
@@ -182,6 +183,7 @@ def test_done_callbacks(self):
 
         callback.assert_called_once_with(mock.sentinel.future)
 
+    @pytest.mark.asyncio
     async def test_metadata(self):
         rpc, call = make_async_rpc()
         bidi_rpc = bidi_async.AsyncBidiRpc(rpc, metadata=mock.sentinel.A)
@@ -192,6 +194,7 @@ async def test_metadata(self):
         rpc.assert_awaited_once()
         assert rpc.call_args.kwargs["metadata"] == mock.sentinel.A
 
+    @pytest.mark.asyncio
     async def test_open(self):
         rpc, call = make_async_rpc()
         bidi_rpc = bidi_async.AsyncBidiRpc(rpc)
@@ -202,6 +205,7 @@ async def test_open(self):
         assert bidi_rpc.is_active
         call.add_done_callback.assert_called_once_with(bidi_rpc._on_call_done)
 
+    @pytest.mark.asyncio
     async def test_open_error_already_open(self):
         rpc, _ = make_async_rpc()
         bidi_rpc = bidi_async.AsyncBidiRpc(rpc)
@@ -211,6 +215,7 @@ async def test_open_error_already_open(self):
         with pytest.raises(ValueError):
             await bidi_rpc.open()
 
+    @pytest.mark.asyncio
     async def test_close(self):
         rpc, call = make_async_rpc()
         bidi_rpc = bidi_async.AsyncBidiRpc(rpc)
@@ -228,10 +233,12 @@ async def test_close(self):
         assert bidi_rpc._initial_request is None
         assert not bidi_rpc._callbacks
 
+    @pytest.mark.asyncio
     async def test_close_no_rpc(self):
         bidi_rpc = bidi_async.AsyncBidiRpc(None)
         await bidi_rpc.close()
 
+    @pytest.mark.asyncio
     async def test_send(self):
         rpc, call = make_async_rpc()
         bidi_rpc = bidi_async.AsyncBidiRpc(rpc)
@@ -242,12 +249,14 @@ async def test_send(self):
         assert bidi_rpc.pending_requests == 1
         assert await bidi_rpc._request_queue.get() is mock.sentinel.request
 
+    @pytest.mark.asyncio
     async def test_send_not_open(self):
         bidi_rpc = bidi_async.AsyncBidiRpc(None)
 
         with pytest.raises(ValueError):
             await bidi_rpc.send(mock.sentinel.request)
 
+    @pytest.mark.asyncio
     async def test_send_dead_rpc(self):
         error = ValueError()
         bidi_rpc = bidi_async.AsyncBidiRpc(None)
@@ -258,15 +267,17 @@ async def test_send_dead_rpc(self):
 
         assert exc_info.value == error
 
+    @pytest.mark.asyncio
     async def test_recv(self):
         bidi_rpc = bidi_async.AsyncBidiRpc(None)
         bidi_rpc.call = mock.create_autospec(aio.Call, instance=True)
-        bidi_rpc.call.read = mock.AsyncMock(return_value=mock.sentinel.response)
+        bidi_rpc.call.read = AsyncMock(return_value=mock.sentinel.response)
 
         response = await bidi_rpc.recv()
 
         assert response == mock.sentinel.response
 
+    @pytest.mark.asyncio
     async def test_recv_not_open(self):
         bidi_rpc = bidi_async.AsyncBidiRpc(None)