feat: add support for asynchronous rest streaming

google/api_core/_rest_streaming_base.py

@@ -27,7 +27,6 @@ class BaseResponseIterator:
     """Base Iterator over REST API responses. This class should not be used directly.
 
     Args:
-        response (requests.Response): An API response object.
         response_message_cls (Union[proto.Message, google.protobuf.message.Message]): A response
         class expected to be returned from an API.
 
@@ -40,8 +39,6 @@ def __init__(
         response_message_cls: Union[proto.Message, google.protobuf.message.Message],
     ):
         self._response_message_cls = response_message_cls
-        # Inner iterator over HTTP response's content.
-        # self._response_itr = self._response.iter_content(decode_unicode=True)
         # Contains a list of JSON responses ready to be sent to user.
         self._ready_objs: Deque[str] = deque()
         # Current JSON response being built.
@@ -100,7 +97,9 @@ def _process_chunk(self, chunk: str):
     def _grab(self):
         # Add extra quotes to make json.loads happy.
         if issubclass(self._response_message_cls, proto.Message):
-            return self._response_message_cls.from_json(self._ready_objs.popleft())
+            return self._response_message_cls.from_json(
+                self._ready_objs.popleft(), ignore_unknown_fields=True
+            )
         elif issubclass(self._response_message_cls, google.protobuf.message.Message):
             return Parse(self._ready_objs.popleft(), self._response_message_cls())
         else:

google/api_core/rest_streaming.py

@@ -31,7 +31,9 @@ class ResponseIterator(BaseResponseIterator):
         class expected to be returned from an API.
 
     Raises:
-        ValueError: If `response_message_cls` is not a subclass of `proto.Message` or `google.protobuf.message.Message`.
+        ValueError:
+            - If `response_message_cls` is not a subclass of `proto.Message` or `google.protobuf.message.Message`.
+            - If `response` is not an instance of `requests.Response`.
     """
 
     def __init__(

google/api_core/rest_streaming_async.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Helpers for server-side asynchronous streaming in REST."""
+"""Helpers for asynchronous server-side streaming in REST."""
 
 from typing import Union
 
@@ -33,7 +33,7 @@ class expected to be returned from an API.
     Raises:
         ValueError:
             - If `response_message_cls` is not a subclass of `proto.Message` or `google.protobuf.message.Message`.
-            - If `response` is not an instance of `google.auth.aio.transport.aiohttp.Response`.
+            - If `response` is not an instance of a subclass of `google.auth.aio.transport.Response`.
     """
 
     def __init__(

tests/conftest.py → tests/helpers.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+"""Helpers for tests"""
 
 import logging
 from typing import List
@@ -59,11 +60,11 @@ def parse_responses(response_message_cls, responses: List[proto.Message]) -> byt
     # in order to be an actual JSON.
     json_responses = [
         (
-            response_message_cls.to_json(r).strip('"')
+            response_message_cls.to_json(response).strip('"')
             if issubclass(response_message_cls, proto.Message)
-            else MessageToJson(r).strip('"')
+            else MessageToJson(response).strip('"')
         )
-        for r in responses
+        for response in responses
     ]
     logging.info(f"Sending JSON stream: {json_responses}")
     ret_val = "[{}]".format(",".join(json_responses))

tests/unit/test_rest_streaming.py

@@ -27,12 +27,12 @@
 from google.api import http_pb2
 from google.api import httpbody_pb2
 
-from ..conftest import Composer, Song, EchoResponse, parse_responses
+from ..helpers import Composer, Song, EchoResponse, parse_responses
 
 
 __protobuf__ = proto.module(package=__name__)
 SEED = int(time.time())
-logging.info(f"Starting rest streaming tests with random seed: {SEED}")
+logging.info(f"Starting sync rest streaming tests with random seed: {SEED}")
 random.seed(SEED)
 
 
@@ -65,6 +65,9 @@ def __init__(
         self._random_split = random_split
         self._response_message_cls = response_cls
 
+    def _parse_responses(self):
+        return parse_responses(self._response_message_cls, self._responses)
+
     def close(self):
         raise NotImplementedError()
 
@@ -74,9 +77,6 @@ def iter_content(self, *args, **kwargs):
             random_split=self._random_split,
         )
 
-    def _parse_responses(self):
-        return parse_responses(self._response_message_cls, self._responses)
-
 
 @pytest.mark.parametrize(
     "random_split,resp_message_is_proto_plus",

tests/unit/test_rest_streaming_async.py

@@ -27,9 +27,9 @@
 from google.api import httpbody_pb2
 from google.auth.aio.transport import Response
 
-from ..conftest import Composer, Song, EchoResponse, parse_responses
+from ..helpers import Composer, Song, EchoResponse, parse_responses
+
 
-# TODO (ohmayr): check if we need to log.
 __protobuf__ = proto.module(package=__name__)
 SEED = int(time.time())
 logging.info(f"Starting async rest streaming tests with random seed: {SEED}")
@@ -46,21 +46,21 @@ class ResponseMock(Response):
     class _ResponseItr(AsyncIterator[bytes]):
         def __init__(self, _response_bytes: bytes, random_split=False):
             self._responses_bytes = _response_bytes
-            self._i = 0
+            self._idx = 0
             self._random_split = random_split
 
         def __aiter__(self):
             return self
 
         async def __anext__(self):
-            if self._i == len(self._responses_bytes):
+            if self._idx >= len(self._responses_bytes):
                 raise StopAsyncIteration
             if self._random_split:
-                n = random.randint(1, len(self._responses_bytes[self._i :]))
+                n = random.randint(1, len(self._responses_bytes[self._idx :]))
             else:
                 n = 1
-            x = self._responses_bytes[self._i : self._i + n]
-            self._i += n
+            x = self._responses_bytes[self._idx : self._idx + n]
+            self._idx += n
             return x
 
     def __init__(
@@ -73,7 +73,17 @@ def __init__(
         self._random_split = random_split
         self._response_message_cls = response_cls
 
+    def _parse_responses(self):
+        return parse_responses(self._response_message_cls, self._responses)
+
+    @property
+    async def headers(self):
+        raise NotImplementedError()
+
     @property
+    async def status_code(self):
+        raise NotImplementedError()
+
     async def close(self):
         raise NotImplementedError()
 
@@ -87,17 +97,6 @@ async def content(self, chunk_size=None):
     async def read(self):
         raise NotImplementedError()
 
-    @property
-    async def headers(self):
-        raise NotImplementedError()
-
-    @property
-    async def status_code(self):
-        raise NotImplementedError()
-
-    def _parse_responses(self):
-        return parse_responses(self._response_message_cls, self._responses)
-
 
 @pytest.mark.asyncio
 @pytest.mark.parametrize(
@@ -119,10 +118,10 @@ async def test_next_simple(random_split, resp_message_is_proto_plus):
         responses=responses, random_split=random_split, response_cls=response_type
     )
     itr = rest_streaming_async.AsyncResponseIterator(resp, response_type)
-    i = 0
+    idx = 0
     async for response in itr:
-        assert response == responses[i]
-        i += 1
+        assert response == responses[idx]
+        idx += 1
 
 
 @pytest.mark.asyncio
@@ -160,11 +159,11 @@ async def test_next_nested(random_split, resp_message_is_proto_plus):
         responses=responses, random_split=random_split, response_cls=response_type
     )
     itr = rest_streaming_async.AsyncResponseIterator(resp, response_type)
-    i = 0
+    idx = 0
     async for response in itr:
-        assert response == responses[i]
-        i += 1
-    assert i == len(responses)
+        assert response == responses[idx]
+        idx += 1
+    assert idx == len(responses)
 
 
 @pytest.mark.asyncio
@@ -198,11 +197,11 @@ async def test_next_stress(random_split, resp_message_is_proto_plus):
         responses=responses, random_split=random_split, response_cls=response_type
     )
     itr = rest_streaming_async.AsyncResponseIterator(resp, response_type)
-    i = 0
+    idx = 0
     async for response in itr:
-        assert response == responses[i]
-        i += 1
-    assert i == n
+        assert response == responses[idx]
+        idx += 1
+    assert idx == n
 
 
 @pytest.mark.asyncio
@@ -266,12 +265,11 @@ async def test_next_escaped_characters_in_string(
         responses=responses, random_split=random_split, response_cls=response_type
     )
     itr = rest_streaming_async.AsyncResponseIterator(resp, response_type)
-
-    i = 0
+    idx = 0
     async for response in itr:
-        assert response == responses[i]
-        i += 1
-    assert i == len(responses)
+        assert response == responses[idx]
+        idx += 1
+    assert idx == len(responses)
 
 
 @pytest.mark.asyncio