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

Lines changed: 1 addition & 1 deletion
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

Lines changed: 118 additions & 0 deletions
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

Lines changed: 8 additions & 74 deletions
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

Lines changed: 83 additions & 0 deletions
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

Lines changed: 13 additions & 1 deletion
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://git@github.com/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)
pFad - Phonifier reborn

Pfad - The Proxy pFad of © 2024 Garber Painting. All rights reserved.

Note: This service is not intended for secure transactions such as banking, social media, email, or purchasing. Use at your own risk. We assume no liability whatsoever for broken pages.


Alternative Proxies:

Alternative Proxy

pFad Proxy

pFad v3 Proxy

pFad v4 Proxy