feat: add support for asynchronous rest streaming

google/api_core/_rest_streaming_base.py

@@ -0,0 +1,109 @@
+# 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.
+# You may obtain a copy of the License at
+#
+#     http://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.
+
+"""Helpers for server-side streaming in REST."""
+
+from collections import deque
+import string
+from typing import Deque, Union
+
+import proto
+import google.protobuf.message
+from google.protobuf.json_format import Parse
+
+
+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.
+
+    Raises:
+        ValueError: If `response_message_cls` is not a subclass of `proto.Message` or `google.protobuf.message.Message`.
+    """
+
+    def __init__(
+        self,
+        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.
+        self._obj = ""
+        # Keeps track of the nesting level within a JSON object.
+        self._level = 0
+        # Keeps track whether HTTP response is currently sending values
+        # inside of a string value.
+        self._in_string = False
+        # Whether an escape symbol "\" was encountered.
+        self._escape_next = False
+
+    def _process_chunk(self, chunk: str):
+        if self._level == 0:
+            if chunk[0] != "[":
+                raise ValueError(
+                    "Can only parse array of JSON objects, instead got %s" % chunk
+                )
+        for char in chunk:
+            if char == "{":
+                if self._level == 1:
+                    # Level 1 corresponds to the outermost JSON object
+                    # (i.e. the one we care about).
+                    self._obj = ""
+                if not self._in_string:
+                    self._level += 1
+                self._obj += char
+            elif char == "}":
+                self._obj += char
+                if not self._in_string:
+                    self._level -= 1
+                if not self._in_string and self._level == 1:
+                    self._ready_objs.append(self._obj)
+            elif char == '"':
+                # Helps to deal with an escaped quotes inside of a string.
+                if not self._escape_next:
+                    self._in_string = not self._in_string
+                self._obj += char
+            elif char in string.whitespace:
+                if self._in_string:
+                    self._obj += char
+            elif char == "[":
+                if self._level == 0:
+                    self._level += 1
+                else:
+                    self._obj += char
+            elif char == "]":
+                if self._level == 1:
+                    self._level -= 1
+                else:
+                    self._obj += char
+            else:
+                self._obj += char
+            self._escape_next = not self._escape_next if char == "\\" else False
+
+    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())
+        elif issubclass(self._response_message_cls, google.protobuf.message.Message):
+            return Parse(self._ready_objs.popleft(), self._response_message_cls())
+        else:
+            raise ValueError(
+                "Response message class must be a subclass of proto.Message or google.protobuf.message.Message."
+            )

google/api_core/rest_streaming.py

@@ -14,17 +14,15 @@
 
 """Helpers for server-side streaming in REST."""
 
-from collections import deque
-import string
-from typing import Deque, Union
+from typing import Union
 
 import proto
 import requests
 import google.protobuf.message
-from google.protobuf.json_format import Parse
+from google.api_core._rest_streaming_base import BaseResponseIterator
 
 
-class ResponseIterator:
+class ResponseIterator(BaseResponseIterator):
     """Iterator over REST API responses.
 
     Args:
@@ -42,68 +40,16 @@ def __init__(
         response_message_cls: Union[proto.Message, google.protobuf.message.Message],
     ):
         self._response = response
-        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.
-        self._obj = ""
-        # Keeps track of the nesting level within a JSON object.
-        self._level = 0
-        # Keeps track whether HTTP response is currently sending values
-        # inside of a string value.
-        self._in_string = False
-        # Whether an escape symbol "\" was encountered.
-        self._escape_next = False
+        super(ResponseIterator, self).__init__(
+            response_message_cls=response_message_cls
+        )
 
     def cancel(self):
         """Cancel existing streaming operation."""
         self._response.close()
 
-    def _process_chunk(self, chunk: str):
-        if self._level == 0:
-            if chunk[0] != "[":
-                raise ValueError(
-                    "Can only parse array of JSON objects, instead got %s" % chunk
-                )
-        for char in chunk:
-            if char == "{":
-                if self._level == 1:
-                    # Level 1 corresponds to the outermost JSON object
-                    # (i.e. the one we care about).
-                    self._obj = ""
-                if not self._in_string:
-                    self._level += 1
-                self._obj += char
-            elif char == "}":
-                self._obj += char
-                if not self._in_string:
-                    self._level -= 1
-                if not self._in_string and self._level == 1:
-                    self._ready_objs.append(self._obj)
-            elif char == '"':
-                # Helps to deal with an escaped quotes inside of a string.
-                if not self._escape_next:
-                    self._in_string = not self._in_string
-                self._obj += char
-            elif char in string.whitespace:
-                if self._in_string:
-                    self._obj += char
-            elif char == "[":
-                if self._level == 0:
-                    self._level += 1
-                else:
-                    self._obj += char
-            elif char == "]":
-                if self._level == 1:
-                    self._level -= 1
-                else:
-                    self._obj += char
-            else:
-                self._obj += char
-            self._escape_next = not self._escape_next if char == "\\" else False
-
     def __next__(self):
         while not self._ready_objs:
             try:
@@ -115,18 +61,5 @@ def __next__(self):
                 raise e
         return self._grab()
 
-    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(), 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:
-            raise ValueError(
-                "Response message class must be a subclass of proto.Message or google.protobuf.message.Message."
-            )
-
     def __iter__(self):
         return self

google/api_core/rest_streaming_async.py

@@ -0,0 +1,74 @@
+# 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.
+# You may obtain a copy of the License at
+#
+#     http://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.
+
+"""Helpers for server-side asynchronous streaming in REST."""
+
+from typing import Union
+
+import proto
+import google.auth.aio.transport
+import google.protobuf.message
+from google.api_core._rest_streaming_base import BaseResponseIterator
+
+
+class AsyncResponseIterator(BaseResponseIterator):
+    """Asynchronous Iterator over REST API responses.
+
+    Args:
+        response (google.auth.aio.transport.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.
+
+    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`.
+    """
+
+    def __init__(
+        self,
+        response: google.auth.aio.transport.Response,
+        response_message_cls: Union[proto.Message, google.protobuf.message.Message],
+    ):
+        self._response = response
+        self._chunk_size = 1024
+        self._response_itr = self._response.content().__aiter__()
+        super(AsyncResponseIterator, self).__init__(
+            response_message_cls=response_message_cls
+        )
+
+    async def cancel(self):
+        """Cancel existing streaming operation."""
+        await self._response.close()
+
+    async def __anext__(self):
+        while not self._ready_objs:
+            try:
+                chunk = await self._response_itr.__anext__()
+                chunk = chunk.decode("utf-8")
+                self._process_chunk(chunk)
+            except StopAsyncIteration as e:
+                if self._level > 0:
+                    raise ValueError("i Unfinished stream: %s" % self._obj)
+                raise e
+            except ValueError as e:
+                raise e
+        return self._grab()
+
+    def __aiter__(self):
+        return self
+
+    async def __aexit__(self, exc_type, exc, tb):
+        """Cancel existing async streaming operation."""
+        await self._response.close()

tests/conftest.py

@@ -0,0 +1,70 @@
+# 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.
+# You may obtain a copy of the License at
+#
+#     http://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.
+
+
+import logging
+from typing import List
+
+import proto
+
+from google.protobuf import duration_pb2
+from google.protobuf import timestamp_pb2
+from google.protobuf.json_format import MessageToJson
+
+
+class Genre(proto.Enum):
+    GENRE_UNSPECIFIED = 0
+    CLASSICAL = 1
+    JAZZ = 2
+    ROCK = 3
+
+
+class Composer(proto.Message):
+    given_name = proto.Field(proto.STRING, number=1)
+    family_name = proto.Field(proto.STRING, number=2)
+    relateds = proto.RepeatedField(proto.STRING, number=3)
+    indices = proto.MapField(proto.STRING, proto.STRING, number=4)
+
+
+class Song(proto.Message):
+    composer = proto.Field(Composer, number=1)
+    title = proto.Field(proto.STRING, number=2)
+    lyrics = proto.Field(proto.STRING, number=3)
+    year = proto.Field(proto.INT32, number=4)
+    genre = proto.Field(Genre, number=5)
+    is_five_mins_longer = proto.Field(proto.BOOL, number=6)
+    score = proto.Field(proto.DOUBLE, number=7)
+    likes = proto.Field(proto.INT64, number=8)
+    duration = proto.Field(duration_pb2.Duration, number=9)
+    date_added = proto.Field(timestamp_pb2.Timestamp, number=10)
+
+
+class EchoResponse(proto.Message):
+    content = proto.Field(proto.STRING, number=1)
+
+
+def parse_responses(response_message_cls, responses: List[proto.Message]) -> bytes:
+    # json.dumps returns a string surrounded with quotes that need to be stripped
+    # in order to be an actual JSON.
+    json_responses = [
+        (
+            response_message_cls.to_json(r).strip('"')
+            if issubclass(response_message_cls, proto.Message)
+            else MessageToJson(r).strip('"')
+        )
+        for r in responses
+    ]
+    logging.info(f"Sending JSON stream: {json_responses}")
+    ret_val = "[{}]".format(",".join(json_responses))
+    return bytes(ret_val, "utf-8")