Skip to content

Commit 1b7bb6d

Browse files
ohmayrgcf-owl-bot[bot]
ohmayr
and
gcf-owl-bot[bot]
authored
feat: add support for asynchronous rest streaming (#686)
* duplicating file to base * restore original file * duplicate file to async * restore original file * duplicate test file for async * restore test file * feat: add support for asynchronous rest streaming * 🦉 Updates from OwlBot post-processor See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md * fix naming issue * fix import module name * pull auth feature branch * revert setup file * address PR comments * 🦉 Updates from OwlBot post-processor See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md * run black * address PR comments * update nox coverage * address PR comments * fix nox session name in workflow * use https for remote repo * add context manager methods * address PR comments * update auth error versions * update import error --------- Co-authored-by: Owl Bot <gcf-owl-bot[bot]@users.noreply.github.com>
1 parent e542124 commit 1b7bb6d

File tree

8 files changed

+679
-128
lines changed

8 files changed

+679
-128
lines changed

‎.github/workflows/unittest.yml

+1-1
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@ jobs:
1111
runs-on: ubuntu-latest
1212
strategy:
1313
matrix:
14-
option: ["", "_grpc_gcp", "_wo_grpc", "_with_prerelease_deps"]
14+
option: ["", "_grpc_gcp", "_wo_grpc", "_with_prerelease_deps", "_with_auth_aio"]
1515
python:
1616
- "3.7"
1717
- "3.8"

‎google/api_core/_rest_streaming_base.py

+118
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,118 @@
1+
# Copyright 2024 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
"""Helpers for server-side streaming in REST."""
16+
17+
from collections import deque
18+
import string
19+
from typing import Deque, Union
20+
import types
21+
22+
import proto
23+
import google.protobuf.message
24+
from google.protobuf.json_format import Parse
25+
26+
27+
class BaseResponseIterator:
28+
"""Base Iterator over REST API responses. This class should not be used directly.
29+
30+
Args:
31+
response_message_cls (Union[proto.Message, google.protobuf.message.Message]): A response
32+
class expected to be returned from an API.
33+
34+
Raises:
35+
ValueError: If `response_message_cls` is not a subclass of `proto.Message` or `google.protobuf.message.Message`.
36+
"""
37+
38+
def __init__(
39+
self,
40+
response_message_cls: Union[proto.Message, google.protobuf.message.Message],
41+
):
42+
self._response_message_cls = response_message_cls
43+
# Contains a list of JSON responses ready to be sent to user.
44+
self._ready_objs: Deque[str] = deque()
45+
# Current JSON response being built.
46+
self._obj = ""
47+
# Keeps track of the nesting level within a JSON object.
48+
self._level = 0
49+
# Keeps track whether HTTP response is currently sending values
50+
# inside of a string value.
51+
self._in_string = False
52+
# Whether an escape symbol "\" was encountered.
53+
self._escape_next = False
54+
55+
self._grab = types.MethodType(self._create_grab(), self)
56+
57+
def _process_chunk(self, chunk: str):
58+
if self._level == 0:
59+
if chunk[0] != "[":
60+
raise ValueError(
61+
"Can only parse array of JSON objects, instead got %s" % chunk
62+
)
63+
for char in chunk:
64+
if char == "{":
65+
if self._level == 1:
66+
# Level 1 corresponds to the outermost JSON object
67+
# (i.e. the one we care about).
68+
self._obj = ""
69+
if not self._in_string:
70+
self._level += 1
71+
self._obj += char
72+
elif char == "}":
73+
self._obj += char
74+
if not self._in_string:
75+
self._level -= 1
76+
if not self._in_string and self._level == 1:
77+
self._ready_objs.append(self._obj)
78+
elif char == '"':
79+
# Helps to deal with an escaped quotes inside of a string.
80+
if not self._escape_next:
81+
self._in_string = not self._in_string
82+
self._obj += char
83+
elif char in string.whitespace:
84+
if self._in_string:
85+
self._obj += char
86+
elif char == "[":
87+
if self._level == 0:
88+
self._level += 1
89+
else:
90+
self._obj += char
91+
elif char == "]":
92+
if self._level == 1:
93+
self._level -= 1
94+
else:
95+
self._obj += char
96+
else:
97+
self._obj += char
98+
self._escape_next = not self._escape_next if char == "\\" else False
99+
100+
def _create_grab(self):
101+
if issubclass(self._response_message_cls, proto.Message):
102+
103+
def grab(this):
104+
return this._response_message_cls.from_json(
105+
this._ready_objs.popleft(), ignore_unknown_fields=True
106+
)
107+
108+
return grab
109+
elif issubclass(self._response_message_cls, google.protobuf.message.Message):
110+
111+
def grab(this):
112+
return Parse(this._ready_objs.popleft(), this._response_message_cls())
113+
114+
return grab
115+
else:
116+
raise ValueError(
117+
"Response message class must be a subclass of proto.Message or google.protobuf.message.Message."
118+
)

‎google/api_core/rest_streaming.py

+8-74
Original file line numberDiff line numberDiff line change
@@ -14,17 +14,15 @@
1414

1515
"""Helpers for server-side streaming in REST."""
1616

17-
from collections import deque
18-
import string
19-
from typing import Deque, Union
17+
from typing import Union
2018

2119
import proto
2220
import requests
2321
import google.protobuf.message
24-
from google.protobuf.json_format import Parse
22+
from google.api_core._rest_streaming_base import BaseResponseIterator
2523

2624

27-
class ResponseIterator:
25+
class ResponseIterator(BaseResponseIterator):
2826
"""Iterator over REST API responses.
2927
3028
Args:
@@ -33,7 +31,8 @@ class ResponseIterator:
3331
class expected to be returned from an API.
3432
3533
Raises:
36-
ValueError: If `response_message_cls` is not a subclass of `proto.Message` or `google.protobuf.message.Message`.
34+
ValueError:
35+
- If `response_message_cls` is not a subclass of `proto.Message` or `google.protobuf.message.Message`.
3736
"""
3837

3938
def __init__(
@@ -42,68 +41,16 @@ def __init__(
4241
response_message_cls: Union[proto.Message, google.protobuf.message.Message],
4342
):
4443
self._response = response
45-
self._response_message_cls = response_message_cls
4644
# Inner iterator over HTTP response's content.
4745
self._response_itr = self._response.iter_content(decode_unicode=True)
48-
# Contains a list of JSON responses ready to be sent to user.
49-
self._ready_objs: Deque[str] = deque()
50-
# Current JSON response being built.
51-
self._obj = ""
52-
# Keeps track of the nesting level within a JSON object.
53-
self._level = 0
54-
# Keeps track whether HTTP response is currently sending values
55-
# inside of a string value.
56-
self._in_string = False
57-
# Whether an escape symbol "\" was encountered.
58-
self._escape_next = False
46+
super(ResponseIterator, self).__init__(
47+
response_message_cls=response_message_cls
48+
)
5949

6050
def cancel(self):
6151
"""Cancel existing streaming operation."""
6252
self._response.close()
6353

64-
def _process_chunk(self, chunk: str):
65-
if self._level == 0:
66-
if chunk[0] != "[":
67-
raise ValueError(
68-
"Can only parse array of JSON objects, instead got %s" % chunk
69-
)
70-
for char in chunk:
71-
if char == "{":
72-
if self._level == 1:
73-
# Level 1 corresponds to the outermost JSON object
74-
# (i.e. the one we care about).
75-
self._obj = ""
76-
if not self._in_string:
77-
self._level += 1
78-
self._obj += char
79-
elif char == "}":
80-
self._obj += char
81-
if not self._in_string:
82-
self._level -= 1
83-
if not self._in_string and self._level == 1:
84-
self._ready_objs.append(self._obj)
85-
elif char == '"':
86-
# Helps to deal with an escaped quotes inside of a string.
87-
if not self._escape_next:
88-
self._in_string = not self._in_string
89-
self._obj += char
90-
elif char in string.whitespace:
91-
if self._in_string:
92-
self._obj += char
93-
elif char == "[":
94-
if self._level == 0:
95-
self._level += 1
96-
else:
97-
self._obj += char
98-
elif char == "]":
99-
if self._level == 1:
100-
self._level -= 1
101-
else:
102-
self._obj += char
103-
else:
104-
self._obj += char
105-
self._escape_next = not self._escape_next if char == "\\" else False
106-
10754
def __next__(self):
10855
while not self._ready_objs:
10956
try:
@@ -115,18 +62,5 @@ def __next__(self):
11562
raise e
11663
return self._grab()
11764

118-
def _grab(self):
119-
# Add extra quotes to make json.loads happy.
120-
if issubclass(self._response_message_cls, proto.Message):
121-
return self._response_message_cls.from_json(
122-
self._ready_objs.popleft(), ignore_unknown_fields=True
123-
)
124-
elif issubclass(self._response_message_cls, google.protobuf.message.Message):
125-
return Parse(self._ready_objs.popleft(), self._response_message_cls())
126-
else:
127-
raise ValueError(
128-
"Response message class must be a subclass of proto.Message or google.protobuf.message.Message."
129-
)
130-
13165
def __iter__(self):
13266
return self

‎google/api_core/rest_streaming_async.py

+83
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,83 @@
1+
# Copyright 2024 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
"""Helpers for asynchronous server-side streaming in REST."""
16+
17+
from typing import Union
18+
19+
import proto
20+
21+
try:
22+
import google.auth.aio.transport
23+
except ImportError as e: # pragma: NO COVER
24+
raise ImportError(
25+
"google-auth>=2.35.0 is required to use asynchronous rest streaming."
26+
) from e
27+
28+
import google.protobuf.message
29+
from google.api_core._rest_streaming_base import BaseResponseIterator
30+
31+
32+
class AsyncResponseIterator(BaseResponseIterator):
33+
"""Asynchronous Iterator over REST API responses.
34+
35+
Args:
36+
response (google.auth.aio.transport.Response): An API response object.
37+
response_message_cls (Union[proto.Message, google.protobuf.message.Message]): A response
38+
class expected to be returned from an API.
39+
40+
Raises:
41+
ValueError:
42+
- If `response_message_cls` is not a subclass of `proto.Message` or `google.protobuf.message.Message`.
43+
"""
44+
45+
def __init__(
46+
self,
47+
response: google.auth.aio.transport.Response,
48+
response_message_cls: Union[proto.Message, google.protobuf.message.Message],
49+
):
50+
self._response = response
51+
self._chunk_size = 1024
52+
self._response_itr = self._response.content().__aiter__()
53+
super(AsyncResponseIterator, self).__init__(
54+
response_message_cls=response_message_cls
55+
)
56+
57+
async def __aenter__(self):
58+
return self
59+
60+
async def cancel(self):
61+
"""Cancel existing streaming operation."""
62+
await self._response.close()
63+
64+
async def __anext__(self):
65+
while not self._ready_objs:
66+
try:
67+
chunk = await self._response_itr.__anext__()
68+
chunk = chunk.decode("utf-8")
69+
self._process_chunk(chunk)
70+
except StopAsyncIteration as e:
71+
if self._level > 0:
72+
raise ValueError("i Unfinished stream: %s" % self._obj)
73+
raise e
74+
except ValueError as e:
75+
raise e
76+
return self._grab()
77+
78+
def __aiter__(self):
79+
return self
80+
81+
async def __aexit__(self, exc_type, exc, tb):
82+
"""Cancel existing async streaming operation."""
83+
await self._response.close()

‎noxfile.py

+13-1
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,7 @@
3838
"unit",
3939
"unit_grpc_gcp",
4040
"unit_wo_grpc",
41+
"unit_with_auth_aio",
4142
"cover",
4243
"pytype",
4344
"mypy",
@@ -109,7 +110,7 @@ def install_prerelease_dependencies(session, constraints_path):
109110
session.install(*other_deps)
110111

111112

112-
def default(session, install_grpc=True, prerelease=False):
113+
def default(session, install_grpc=True, prerelease=False, install_auth_aio=False):
113114
"""Default unit test session.
114115
115116
This is intended to be run **without** an interpreter set, so
@@ -144,6 +145,11 @@ def default(session, install_grpc=True, prerelease=False):
144145
f"{constraints_dir}/constraints-{session.python}.txt",
145146
)
146147

148+
if install_auth_aio:
149+
session.install(
150+
"google-auth @ git+https://[email protected]/googleapis/google-auth-library-python@8833ad6f92c3300d6645355994c7db2356bd30ad"
151+
)
152+
147153
# Print out package versions of dependencies
148154
session.run(
149155
"python", "-c", "import google.protobuf; print(google.protobuf.__version__)"
@@ -229,6 +235,12 @@ def unit_wo_grpc(session):
229235
default(session, install_grpc=False)
230236

231237

238+
@nox.session(python=PYTHON_VERSIONS)
239+
def unit_with_auth_aio(session):
240+
"""Run the unit test suite with google.auth.aio installed"""
241+
default(session, install_auth_aio=True)
242+
243+
232244
@nox.session(python=DEFAULT_PYTHON_VERSION)
233245
def lint_setup_py(session):
234246
"""Verify that setup.py is valid (including RST check)."""

0 commit comments

Comments
 (0)