From 800f1f77171218bec1e5216836e95cca3a6c4724 Mon Sep 17 00:00:00 2001 From: Antonio Aranda <102337110+arandito@users.noreply.github.com> Date: Fri, 21 Aug 2026 10:28:33 -0400 Subject: [PATCH] smithy-http: stream aiohttp response bodies --- ...gfix-b0c1d627a5c34c5b953a785d049c0bb6.json | 4 + .../src/smithy_http/aio/aiohttp.py | 35 ++++-- .../tests/unit/aio/test_aiohttp.py | 106 +++++++++++++++--- 3 files changed, 119 insertions(+), 26 deletions(-) create mode 100644 packages/smithy-http/.changes/next-release/smithy-http-bugfix-b0c1d627a5c34c5b953a785d049c0bb6.json diff --git a/packages/smithy-http/.changes/next-release/smithy-http-bugfix-b0c1d627a5c34c5b953a785d049c0bb6.json b/packages/smithy-http/.changes/next-release/smithy-http-bugfix-b0c1d627a5c34c5b953a785d049c0bb6.json new file mode 100644 index 000000000..9a71763ec --- /dev/null +++ b/packages/smithy-http/.changes/next-release/smithy-http-bugfix-b0c1d627a5c34c5b953a785d049c0bb6.json @@ -0,0 +1,4 @@ +{ + "type": "bugfix", + "description": "Stream aiohttp response bodies incrementally to enable true output event streams." +} \ No newline at end of file diff --git a/packages/smithy-http/src/smithy_http/aio/aiohttp.py b/packages/smithy-http/src/smithy_http/aio/aiohttp.py index b2495c118..3f4266cd6 100644 --- a/packages/smithy-http/src/smithy_http/aio/aiohttp.py +++ b/packages/smithy-http/src/smithy_http/aio/aiohttp.py @@ -1,5 +1,6 @@ # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. # SPDX-License-Identifier: Apache-2.0 +from collections.abc import AsyncIterable, AsyncIterator from copy import copy, deepcopy from itertools import chain from typing import TYPE_CHECKING, Any, Self @@ -22,7 +23,6 @@ from smithy_core.aio.interfaces import StreamingBlob from smithy_core.aio.types import AsyncBytesReader -from smithy_core.aio.utils import async_list from smithy_core.exceptions import MissingDependencyError from smithy_core.interfaces import URI @@ -49,13 +49,28 @@ def __post_init__(self) -> None: _assert_aiohttp() +class _AIOHTTPStreamingBody(AsyncIterable[bytes]): + """Streams a response body, releasing the response once it is done.""" + + def __init__(self, response: "aiohttp.ClientResponse") -> None: + self._response = response + + def __aiter__(self) -> AsyncIterator[bytes]: + # The reader iterates by line, so use iter_any to get raw chunks. + return self._response.content.iter_any() + + async def close(self) -> None: + # Pools the connection if the body was read to completion, closes it if not. + self._response.release() + + class AIOHTTPClient(HTTPClient): """Implementation of :py:class:`.interfaces.HTTPClient` using aiohttp.""" TIMEOUT_EXCEPTIONS = (TimeoutError,) - # aiohttp has no HTTP/2 support and this client fully buffers the response - # before returning, so it can never interleave request and response data. + # aiohttp has no HTTP/2 support, so it can never interleave request and + # response data. SUPPORTS_DUPLEX_STREAMING = False def __init__( @@ -113,15 +128,19 @@ async def send( # The typing on `params` is incorrect, it'll happily accept a mapping whose # values are lists (or tuples) and produce expected values. # See: https://github.com/aio-libs/aiohttp/issues/8563 - async with self._session.request( + resp = await self._session.request( method=request.method, url=self._serialize_uri_without_query(request.destination), params=parse_qs(request.destination.query), # type: ignore headers=headers_list, data=body, allow_redirects=False, - ) as resp: - return await self._marshal_response(resp) + ) + try: + return self._marshal_response(resp) + except BaseException: + resp.release() + raise async def close(self) -> None: """Close the underlying aiohttp session and its connection pool.""" @@ -163,7 +182,7 @@ def _serialize_uri_without_query(self, uri: URI) -> yarl.URL: encoded=True, ) - async def _marshal_response( + def _marshal_response( self, aiohttp_resp: "aiohttp.ClientResponse" ) -> HTTPResponseInterface: """Convert a ``aiohttp.ClientResponse`` to a ``smithy_http.aio.HTTPResponse``""" @@ -181,7 +200,7 @@ async def _marshal_response( return HTTPResponse( status=aiohttp_resp.status, fields=headers, - body=async_list([await aiohttp_resp.read()]), + body=_AIOHTTPStreamingBody(aiohttp_resp), reason=aiohttp_resp.reason, ) diff --git a/packages/smithy-http/tests/unit/aio/test_aiohttp.py b/packages/smithy-http/tests/unit/aio/test_aiohttp.py index c7a9beb25..941215f18 100644 --- a/packages/smithy-http/tests/unit/aio/test_aiohttp.py +++ b/packages/smithy-http/tests/unit/aio/test_aiohttp.py @@ -3,14 +3,14 @@ # pyright: reportPrivateUsage=false from collections.abc import AsyncIterator from typing import Any, cast -from unittest.mock import AsyncMock, MagicMock +from unittest.mock import AsyncMock, MagicMock, patch import pytest from smithy_core import URI from smithy_core.aio.types import AsyncBytesReader from smithy_http import Field, Fields from smithy_http.aio import HTTPRequest -from smithy_http.aio.aiohttp import AIOHTTPClient +from smithy_http.aio.aiohttp import AIOHTTPClient, _AIOHTTPStreamingBody from smithy_http.exceptions import SmithyHTTPError @@ -20,11 +20,19 @@ def _create_client() -> tuple[AIOHTTPClient, MagicMock]: session = MagicMock() session.close = AsyncMock() - session.request.return_value.__aenter__ = AsyncMock(return_value=response) - session.request.return_value.__aexit__ = AsyncMock(return_value=None) + session.request = AsyncMock(return_value=response) return AIOHTTPClient(_session=cast(Any, session)), session +def _create_request() -> HTTPRequest: + return HTTPRequest( + method="GET", + destination=URI(scheme="https", host="example.com", path="/"), + body=AsyncBytesReader(b""), + fields=Fields(), + ) + + async def test_close_closes_session() -> None: client, session = _create_client() @@ -57,14 +65,8 @@ async def test_context_manager_closes_session() -> None: async def test_send_omits_empty_async_reader_body() -> None: client, session = _create_client() - request = HTTPRequest( - method="GET", - destination=URI(scheme="https", host="example.com", path="/"), - body=AsyncBytesReader(b""), - fields=Fields(), - ) - await client.send(request) + await client.send(_create_request()) assert session.request.call_args.kwargs["data"] is None @@ -86,18 +88,86 @@ async def test_send_preserves_explicitly_framed_empty_body() -> None: async def test_send_disables_redirects() -> None: client, session = _create_client() - request = HTTPRequest( - method="GET", - destination=URI(scheme="https", host="example.com", path="/"), - body=AsyncBytesReader(b""), - fields=Fields(), - ) - await client.send(request) + await client.send(_create_request()) assert session.request.call_args.kwargs["allow_redirects"] is False +async def test_send_streams_response_body_and_releases_it_on_close() -> None: + async def chunks() -> AsyncIterator[bytes]: + yield b"first" + yield b"second" + + client, session = _create_client() + aiohttp_response = session.request.return_value + aiohttp_response.content.iter_any.return_value = chunks() + + response = await client.send(_create_request()) + + aiohttp_response.read.assert_not_awaited() + aiohttp_response.content.iter_any.assert_not_called() + assert isinstance(response.body, _AIOHTTPStreamingBody) + assert [chunk async for chunk in response.body] == [ + b"first", + b"second", + ] + aiohttp_response.content.iter_any.assert_called_once_with() + aiohttp_response.release.assert_not_called() + + await response.body.close() + aiohttp_response.release.assert_called_once_with() + + +async def test_non_streaming_response_body_can_be_consumed() -> None: + async def chunks() -> AsyncIterator[bytes]: + yield b'{"message":' + yield b'"hello"}' + + client, session = _create_client() + aiohttp_response = session.request.return_value + aiohttp_response.content.iter_any.return_value = chunks() + + response = await client.send(_create_request()) + + assert await response.consume_body_async() == b'{"message":"hello"}' + aiohttp_response.content.iter_any.assert_called_once_with() + + +async def test_response_body_close_releases_partially_consumed_response() -> None: + async def chunks() -> AsyncIterator[bytes]: + yield b"first" + yield b"second" + + client, session = _create_client() + aiohttp_response = session.request.return_value + aiohttp_response.content.iter_any.return_value = chunks() + + response = await client.send(_create_request()) + assert isinstance(response.body, _AIOHTTPStreamingBody) + body_iterator = aiter(response.body) + assert await anext(body_iterator) == b"first" + + await response.body.close() + + aiohttp_response.release.assert_called_once_with() + + +async def test_send_releases_response_when_marshaling_fails() -> None: + client, session = _create_client() + aiohttp_response = session.request.return_value + + with ( + patch.object( + client, "_marshal_response", side_effect=ValueError("invalid response") + ), + pytest.raises(ValueError, match="invalid response"), + ): + await client.send(_create_request()) + + aiohttp_response.release.assert_called_once_with() + + async def test_prepare_body_preserves_nonempty_reader_position() -> None: client, _ = _create_client() body = AsyncBytesReader(b"request body")