Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: add support for asynchronous rest streaming #686

Merged
merged 30 commits into from
Sep 18, 2024
Merged
Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
30 commits
Select commit Hold shift + click to select a range
4df2120
duplicating file to base
ohmayr Aug 23, 2024
26f52a5
restore original file
ohmayr Aug 23, 2024
204920a
duplicate file to async
ohmayr Aug 23, 2024
2014ef6
restore original file
ohmayr Aug 23, 2024
fdeb437
Merge branch 'dups' into add-support-for-async-rest-streaming
ohmayr Aug 23, 2024
8014f47
duplicate test file for async
ohmayr Aug 23, 2024
e84f03a
restore test file
ohmayr Aug 23, 2024
91f5c13
Merge branch 'dups' into add-support-for-async-rest-streaming
ohmayr Aug 23, 2024
0c2ee5c
feat: add support for asynchronous rest streaming
ohmayr Aug 23, 2024
75ae7d7
🦉 Updates from OwlBot post-processor
gcf-owl-bot[bot] Aug 23, 2024
db7cfb3
fix naming issue
ohmayr Aug 23, 2024
9da0113
Merge branch 'add-support-for-async-rest-streaming' of github.com:goo…
ohmayr Aug 23, 2024
87eeca3
fix import module name
ohmayr Aug 23, 2024
d6abddd
pull auth feature branch
ohmayr Aug 23, 2024
a6a648d
revert setup file
ohmayr Aug 23, 2024
8d30b2d
address PR comments
ohmayr Aug 24, 2024
0b51b09
🦉 Updates from OwlBot post-processor
gcf-owl-bot[bot] Aug 24, 2024
7d8f1e1
run black
ohmayr Aug 24, 2024
d793069
address PR comments
ohmayr Aug 27, 2024
5f48f72
Merge branch 'add-support-for-async-rest-streaming' of github.com:goo…
ohmayr Aug 27, 2024
70d1cdb
update nox coverage
ohmayr Aug 27, 2024
ee3a5ac
address PR comments
ohmayr Aug 27, 2024
5a6e488
fix nox session name in workflow
ohmayr Aug 27, 2024
f2180da
use https for remote repo
ohmayr Aug 27, 2024
bc811e5
add context manager methods
ohmayr Aug 27, 2024
fc74ae0
address PR comments
ohmayr Sep 11, 2024
1603818
Merge branch 'main' into add-support-for-async-rest-streaming
ohmayr Sep 11, 2024
cee06ab
update auth error versions
ohmayr Sep 18, 2024
380dedf
Merge branch 'main' into add-support-for-async-rest-streaming
ohmayr Sep 18, 2024
750d4cb
update import error
ohmayr Sep 18, 2024
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 3 additions & 4 deletions google/api_core/_rest_streaming_base.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.

Expand All @@ -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.
Expand Down Expand Up @@ -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:
Expand Down
4 changes: 3 additions & 1 deletion google/api_core/rest_streaming.py
Original file line number Diff line number Diff line change
Expand Up @@ -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`.
ohmayr marked this conversation as resolved.
Show resolved Hide resolved
- If `response` is not an instance of `requests.Response`.
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Where is this exception raised? Line 46 will succeed with a matching method of any class, right?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think it's reasonable to not raise an error for this. We've added the relevant type hints. Cleaned up the docstring.

"""

def __init__(
Expand Down
4 changes: 2 additions & 2 deletions google/api_core/rest_streaming_async.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand All @@ -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__(
Expand Down
7 changes: 4 additions & 3 deletions tests/conftest.py → tests/helpers.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
ohmayr marked this conversation as resolved.
Show resolved Hide resolved
]
logging.info(f"Sending JSON stream: {json_responses}")
ret_val = "[{}]".format(",".join(json_responses))
Expand Down
10 changes: 5 additions & 5 deletions tests/unit/test_rest_streaming.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
vchudnov-g marked this conversation as resolved.
Show resolved Hide resolved


Expand Down Expand Up @@ -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()

Expand All @@ -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",
Expand Down
66 changes: 32 additions & 34 deletions tests/unit/test_rest_streaming_async.py
Original file line number Diff line number Diff line change
Expand Up @@ -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}")
Expand All @@ -46,21 +46,21 @@ class ResponseMock(Response):
class _ResponseItr(AsyncIterator[bytes]):
def __init__(self, _response_bytes: bytes, random_split=False):
ohmayr marked this conversation as resolved.
Show resolved Hide resolved
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__(
Expand All @@ -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
vchudnov-g marked this conversation as resolved.
Show resolved Hide resolved
async def status_code(self):
raise NotImplementedError()

async def close(self):
raise NotImplementedError()

Expand All @@ -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(
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down
Loading