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.

Sign up for GitHub

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

Jump to bottom

fix: store prefetched exceptions #733

Closed
wants to merge 4 commits into from
Closed

fix: store prefetched exceptions #733

Show file tree
Hide file tree
Changes from 3 commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
d648402
fix: store prefetched exceptions
daniel-sanche Oct 9, 2024
30ded2c
fixed tests
daniel-sanche Oct 9, 2024
bfea5ae
move stored error into own variable
daniel-sanche Oct 9, 2024
0708409
optimize by using extra flag
daniel-sanche Oct 9, 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: 7 additions & 0 deletions google/api_core/grpc_helpers.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -97,6 +97,9 @@ def __init__(self, wrapped, prefetch_first_result=True):
except StopIteration:
# ignore stop iteration at this time. This should be handled outside of retry.
pass
except grpc.RpcError as exc:
# If the pre-fetch fails, store exception to be raised on next() call.
self._stored_first_error = exc

def __iter__(self) -> Iterator[P]:
"""This iterator is also an iterable that returns itself."""
Expand All @@ -113,6 +116,10 @@ def __next__(self) -> P:
result = self._stored_first_result
del self._stored_first_result
return result
elif hasattr(self, "_stored_first_error"):
exc = self._stored_first_error
del self._stored_first_error
raise exc
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'm considering whether it's worth refactoring both these variables, so we don't have to do two hasattr lookups for each stream result. For now I'm just following the existing code though

return next(self._wrapped)
except grpc.RpcError as exc:
# If the stream has already returned data, we cannot recover here.
Expand Down
20 changes: 14 additions & 6 deletions tests/unit/test_grpc_helpers.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -96,11 +96,17 @@ def test_ctor_explicit(self):
assert list(wrapped) == ["a", "b", "c"]

def test_ctor_w_rpc_error_on_prefetch(self):
"""
prefetch errors should not be raised until the first iteration.
"""
wrapped = mock.MagicMock()
wrapped.__next__.side_effect = grpc.RpcError()
exc = grpc.RpcError()
wrapped.__next__.side_effect = exc

with pytest.raises(grpc.RpcError):
self._make_one(wrapped)
obj = self._make_one(wrapped)
assert obj._stored_first_error is exc
with pytest.raises(exceptions.GoogleAPICallError):
next(obj)

def test___iter__(self):
wrapped = self._make_wrapped("a", "b", "c")
Expand Down Expand Up @@ -323,11 +329,13 @@ def test_wrap_stream_errors_iterator_initialization():

wrapped_callable = grpc_helpers._wrap_stream_errors(callable_)

with pytest.raises(exceptions.ServiceUnavailable) as exc_info:
wrapped_callable(1, 2, three="four")
it = wrapped_callable(1, 2, three="four")
assert it._stored_first_error is grpc_error

callable_.assert_called_once_with(1, 2, three="four")
assert exc_info.value.response == grpc_error
# should raise on first iteration
with pytest.raises(exceptions.ServiceUnavailable):
next(it)


def test_wrap_stream_errors_during_iteration():
Expand Down