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..c08007c8f 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,7 @@ 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.aio.utils import read_streaming_blob_async from smithy_core.exceptions import MissingDependencyError from smithy_core.interfaces import URI @@ -49,13 +50,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__( @@ -102,26 +118,39 @@ async def send( ) body: StreamingBlob | None = request.body - if ( - "content-length" not in request.fields - and "transfer-encoding" not in request.fields - ): + if "transfer-encoding" in request.fields: + # The caller explicitly opted into streamed (chunked) framing. + if not isinstance(body, AsyncBytesReader): + body = AsyncBytesReader(body) + elif "content-length" in request.fields: + # The request was signed with this Content-Length. Handing aiohttp an + # async iterable of unknown size makes it fall back to chunked + # transfer encoding, which drops the signed Content-Length (and + # injects Content-Type: application/octet-stream). Both are signed + # headers, so the transmitted request no longer matches the SigV4 + # signature, yielding 403 InvalidSignatureException (or 415 at + # services that validate media type first). Buffer to a fixed-length + # bytes payload so aiohttp preserves the framing that was signed. + body = await read_streaming_blob_async(body) + else: body = await self._prepare_body(body) - elif not isinstance(body, AsyncBytesReader): - body = AsyncBytesReader(body) # 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 +192,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 +210,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..21b1a9524 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,47 +65,143 @@ async def test_context_manager_closes_session() -> None: async def test_send_omits_empty_async_reader_body() -> None: client, session = _create_client() + + await client.send(_create_request()) + + assert session.request.call_args.kwargs["data"] is None + + +async def test_send_buffers_sized_body_to_fixed_length_payload() -> None: + # A request carrying a signed Content-Length must be sent as a fixed-length + # bytes payload, not an unknown-size async iterable. The latter makes aiohttp + # fall back to chunked transfer encoding, which drops the signed + # Content-Length and breaks SigV4 (403) / trips media-type validation (415). + client, session = _create_client() request = HTTPRequest( - method="GET", + method="POST", destination=URI(scheme="https", host="example.com", path="/"), - body=AsyncBytesReader(b""), - fields=Fields(), + body=AsyncBytesReader(b'{"hello":"world"}'), + fields=Fields([Field(name="content-length", values=["17"])]), ) await client.send(request) - assert session.request.call_args.kwargs["data"] is None + assert session.request.call_args.kwargs["data"] == b'{"hello":"world"}' -async def test_send_preserves_explicitly_framed_empty_body() -> None: +async def test_send_buffers_empty_sized_body_to_empty_payload() -> None: client, session = _create_client() - body = AsyncBytesReader(b"") request = HTTPRequest( method="GET", destination=URI(scheme="https", host="example.com", path="/"), - body=body, + body=AsyncBytesReader(b""), fields=Fields([Field(name="content-length", values=["0"])]), ) await client.send(request) - assert session.request.call_args.kwargs["data"] is body + assert session.request.call_args.kwargs["data"] == b"" -async def test_send_disables_redirects() -> None: +async def test_send_streams_body_with_explicit_transfer_encoding() -> None: + # When the caller explicitly opts into chunked framing, the body is streamed + # rather than buffered. client, session = _create_client() + body = AsyncBytesReader(b"streamed body") request = HTTPRequest( - method="GET", + method="POST", destination=URI(scheme="https", host="example.com", path="/"), - body=AsyncBytesReader(b""), - fields=Fields(), + body=body, + fields=Fields([Field(name="transfer-encoding", values=["chunked"])]), ) await client.send(request) + assert session.request.call_args.kwargs["data"] is body + + +async def test_send_disables_redirects() -> None: + client, session = _create_client() + + 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")