From 72659c758b44b49fa6961938028d85741c6fdfd5 Mon Sep 17 00:00:00 2001 From: Sam Bull Date: Sun, 22 Sep 2024 19:03:05 +0100 Subject: [PATCH 01/11] Stop raise_for_status() releasing when in a context --- aiohttp/client.py | 26 ++++------------------- aiohttp/client_reqrep.py | 10 ++++++++- tests/test_client_functional.py | 37 +++++++++++++++++++++++++++++++++ 3 files changed, 50 insertions(+), 23 deletions(-) diff --git a/aiohttp/client.py b/aiohttp/client.py index 26af65cb7e9..bbbe3717294 100644 --- a/aiohttp/client.py +++ b/aiohttp/client.py @@ -1295,11 +1295,7 @@ def __iter__(self) -> Generator[Any, None, _RetType]: async def __aenter__(self) -> _RetType: self._resp = await self._coro - return self._resp - - -class _RequestContextManager(_BaseRequestContextManager[ClientResponse]): - __slots__ = () + return await self._resp.__aenter__() async def __aexit__( self, @@ -1307,25 +1303,11 @@ async def __aexit__( exc: Optional[BaseException], tb: Optional[TracebackType], ) -> None: - # We're basing behavior on the exception as it can be caused by - # user code unrelated to the status of the connection. If you - # would like to close a connection you must do that - # explicitly. Otherwise connection error handling should kick in - # and close/recycle the connection as required. - self._resp.release() - await self._resp.wait_for_close() - + await self._resp.__aexit__(exc_type, exc, tb) -class _WSRequestContextManager(_BaseRequestContextManager[ClientWebSocketResponse]): - __slots__ = () - async def __aexit__( - self, - exc_type: Optional[Type[BaseException]], - exc: Optional[BaseException], - tb: Optional[TracebackType], - ) -> None: - await self._resp.close() +_RequestContextManager = _BaseRequestContextManager[ClientResponse] +_WSRequestContextManager = _BaseRequestContextManager[ClientWebSocketResponse] class _SessionRequestContextManager: diff --git a/aiohttp/client_reqrep.py b/aiohttp/client_reqrep.py index 42d33b0c3cd..319a12a09f1 100644 --- a/aiohttp/client_reqrep.py +++ b/aiohttp/client_reqrep.py @@ -737,6 +737,7 @@ class ClientResponse(HeadersMixin): # post-init stage allows to not change ctor signature _closed = True # to allow __del__ for non-initialized properly response _released = False + _in_context = False __writer = None def __init__( @@ -1022,7 +1023,12 @@ def raise_for_status(self) -> None: if not self.ok: # reason should always be not None for a started response assert self.reason is not None - self.release() + + # If we're in a context we can rely on __aexit__() to release as the + # exception propagates. + if not self._in_context: + self.release() + raise ClientResponseError( self.request_info, self.history, @@ -1144,6 +1150,7 @@ async def json( return loads(self._body.decode(encoding)) # type: ignore[union-attr] async def __aenter__(self) -> "ClientResponse": + self._in_context = True return self async def __aexit__( @@ -1152,6 +1159,7 @@ async def __aexit__( exc_val: Optional[BaseException], exc_tb: Optional[TracebackType], ) -> None: + self._in_context = False # similar to _RequestContextManager, we do not need to check # for exceptions, response object can close connection # if state is broken diff --git a/tests/test_client_functional.py b/tests/test_client_functional.py index 437936e97b8..277bba1f1b8 100644 --- a/tests/test_client_functional.py +++ b/tests/test_client_functional.py @@ -36,6 +36,7 @@ from aiohttp import Fingerprint, ServerFingerprintMismatch, hdrs, web from aiohttp.abc import AbstractResolver, ResolveResult from aiohttp.client_exceptions import ( + ClientResponseError, InvalidURL, InvalidUrlClientError, InvalidUrlRedirectClientError, @@ -3688,6 +3689,42 @@ async def handler(request: web.Request) -> web.Response: await resp.read() +async def test_read_after_catch_raise_for_status(aiohttp_client: AiohttpClient) -> None: + async def handler(request: web.Request) -> web.Response: + return web.Response(body=b"data", status=404) + + app = web.Application() + app.add_routes([web.get("/", handler)]) + + client = await aiohttp_client(app) + + async with client.get("/") as resp: + with pytest.raises(ClientResponseError): + # Should not release response when in async with context. + await resp.raise_for_status() + + result = await resp.read() + assert result == b"data" + + +async def test_read_after_raise_outside_context(aiohttp_client: AiohttpClient) -> None: + async def handler(request: web.Request) -> web.Response: + return web.Response(body=b"data", status=404) + + app = web.Application() + app.add_routes([web.get("/", handler)]) + + client = await aiohttp_client(app) + + resp = await client.get("/") + with pytest.raises(ClientResponseError): + # No async with, so should release and therefore read() will fail. + await resp.raise_for_status() + + with pytest.raises(aiohttp.ClientConnectionError): + result = await resp.read() + + async def test_read_from_closed_content(aiohttp_client: AiohttpClient) -> None: async def handler(request: web.Request) -> web.Response: return web.Response(body=b"data") From c80785c000ea4771a200c54b9cf4affaabb9950d Mon Sep 17 00:00:00 2001 From: Sam Bull Date: Sun, 22 Sep 2024 19:05:04 +0100 Subject: [PATCH 02/11] Create 9239.bugfix.rst --- CHANGES/9239.bugfix.rst | 1 + 1 file changed, 1 insertion(+) create mode 100644 CHANGES/9239.bugfix.rst diff --git a/CHANGES/9239.bugfix.rst b/CHANGES/9239.bugfix.rst new file mode 100644 index 00000000000..20533af7de1 --- /dev/null +++ b/CHANGES/9239.bugfix.rst @@ -0,0 +1 @@ +Stopped ``ClientResponse.raise_for_status()`` releasing the connection within an `async with` context -- by :user:`Dreamsorcerer`. From 77fac64c5dcb2ff1b6cc316b92de1d1eaae45230 Mon Sep 17 00:00:00 2001 From: Sam Bull Date: Sun, 22 Sep 2024 19:06:40 +0100 Subject: [PATCH 03/11] Update tests/test_client_functional.py --- tests/test_client_functional.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_client_functional.py b/tests/test_client_functional.py index 277bba1f1b8..da098c50950 100644 --- a/tests/test_client_functional.py +++ b/tests/test_client_functional.py @@ -3722,7 +3722,7 @@ async def handler(request: web.Request) -> web.Response: await resp.raise_for_status() with pytest.raises(aiohttp.ClientConnectionError): - result = await resp.read() + await resp.read() async def test_read_from_closed_content(aiohttp_client: AiohttpClient) -> None: From f0de714448911ea50d4ce4cc7c0820ef8ad7829f Mon Sep 17 00:00:00 2001 From: Sam Bull Date: Sun, 22 Sep 2024 19:28:47 +0100 Subject: [PATCH 04/11] Update CHANGES/9239.bugfix.rst MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: Sviatoslav Sydorenko (Святослав Сидоренко) --- CHANGES/9239.bugfix.rst | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CHANGES/9239.bugfix.rst b/CHANGES/9239.bugfix.rst index 20533af7de1..c2225a01a2b 100644 --- a/CHANGES/9239.bugfix.rst +++ b/CHANGES/9239.bugfix.rst @@ -1 +1 @@ -Stopped ``ClientResponse.raise_for_status()`` releasing the connection within an `async with` context -- by :user:`Dreamsorcerer`. +Stopped ``ClientResponse.raise_for_status()`` releasing the connection within an ``async with`` context -- by :user:`Dreamsorcerer`. From 609d5f55a2810f480beb9d05e721cfa357a09517 Mon Sep 17 00:00:00 2001 From: Sam Bull Date: Sun, 22 Sep 2024 19:31:57 +0100 Subject: [PATCH 05/11] Update CHANGES/9239.bugfix.rst MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: Sviatoslav Sydorenko (Святослав Сидоренко) --- CHANGES/9239.bugfix.rst | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CHANGES/9239.bugfix.rst b/CHANGES/9239.bugfix.rst index c2225a01a2b..bfddfc3a07e 100644 --- a/CHANGES/9239.bugfix.rst +++ b/CHANGES/9239.bugfix.rst @@ -1 +1 @@ -Stopped ``ClientResponse.raise_for_status()`` releasing the connection within an ``async with`` context -- by :user:`Dreamsorcerer`. +Changed ``ClientResponse.raise_for_status()`` to only release the connection when invoked outside an ``async with`` context -- by :user:`Dreamsorcerer`. From 90a56dca05027d919762b53f79b716368f89275d Mon Sep 17 00:00:00 2001 From: Sam Bull Date: Sun, 22 Sep 2024 19:41:55 +0100 Subject: [PATCH 06/11] Apply suggestions from code review MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: Sviatoslav Sydorenko (Святослав Сидоренко) --- CHANGES/9239.bugfix.rst | 2 +- tests/test_client_functional.py | 6 +++--- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/CHANGES/9239.bugfix.rst b/CHANGES/9239.bugfix.rst index bfddfc3a07e..95b229742ce 100644 --- a/CHANGES/9239.bugfix.rst +++ b/CHANGES/9239.bugfix.rst @@ -1 +1 @@ -Changed ``ClientResponse.raise_for_status()`` to only release the connection when invoked outside an ``async with`` context -- by :user:`Dreamsorcerer`. +Changed :py:meth:`ClientResponse.raise_for_status() ` to only release the connection when invoked outside an ``async with`` context -- by :user:`Dreamsorcerer`. diff --git a/tests/test_client_functional.py b/tests/test_client_functional.py index da098c50950..26c27a31c65 100644 --- a/tests/test_client_functional.py +++ b/tests/test_client_functional.py @@ -3699,7 +3699,7 @@ async def handler(request: web.Request) -> web.Response: client = await aiohttp_client(app) async with client.get("/") as resp: - with pytest.raises(ClientResponseError): + with pytest.raises(ClientResponseError, match="404"): # Should not release response when in async with context. await resp.raise_for_status() @@ -3717,11 +3717,11 @@ async def handler(request: web.Request) -> web.Response: client = await aiohttp_client(app) resp = await client.get("/") - with pytest.raises(ClientResponseError): + with pytest.raises(ClientResponseError, match="404"): # No async with, so should release and therefore read() will fail. await resp.raise_for_status() - with pytest.raises(aiohttp.ClientConnectionError): + with pytest.raises(aiohttp.ClientConnectionError, matches="Connection closed"): await resp.read() From 7b24667ff8b800167b0aaeae89375f3e668fa1e0 Mon Sep 17 00:00:00 2001 From: Sam Bull Date: Sun, 22 Sep 2024 19:42:42 +0100 Subject: [PATCH 07/11] Update tests/test_client_functional.py --- tests/test_client_functional.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_client_functional.py b/tests/test_client_functional.py index 26c27a31c65..1ba02a8b824 100644 --- a/tests/test_client_functional.py +++ b/tests/test_client_functional.py @@ -3721,7 +3721,7 @@ async def handler(request: web.Request) -> web.Response: # No async with, so should release and therefore read() will fail. await resp.raise_for_status() - with pytest.raises(aiohttp.ClientConnectionError, matches="Connection closed"): + with pytest.raises(aiohttp.ClientConnectionError, match="Connection closed"): await resp.read() From f715b63e088d78c5f02fb50aa367df9b588fc5c5 Mon Sep 17 00:00:00 2001 From: Sam Bull Date: Sun, 22 Sep 2024 19:44:27 +0100 Subject: [PATCH 08/11] Update tests/test_client_functional.py MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: Sviatoslav Sydorenko (Святослав Сидоренко) --- tests/test_client_functional.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_client_functional.py b/tests/test_client_functional.py index 1ba02a8b824..c8cfe0e6b62 100644 --- a/tests/test_client_functional.py +++ b/tests/test_client_functional.py @@ -3721,7 +3721,7 @@ async def handler(request: web.Request) -> web.Response: # No async with, so should release and therefore read() will fail. await resp.raise_for_status() - with pytest.raises(aiohttp.ClientConnectionError, match="Connection closed"): + with pytest.raises(aiohttp.ClientConnectionError, match=r"^Connection closed$"): await resp.read() From 4eab6a9d60bb24531fc20f22f141e1f0f4c0ec07 Mon Sep 17 00:00:00 2001 From: Sam Bull Date: Sun, 22 Sep 2024 21:48:15 +0100 Subject: [PATCH 09/11] Update client.py --- aiohttp/client.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/aiohttp/client.py b/aiohttp/client.py index bbbe3717294..39b6b337f0a 100644 --- a/aiohttp/client.py +++ b/aiohttp/client.py @@ -219,7 +219,7 @@ class ClientTimeout: # https://www.rfc-editor.org/rfc/rfc9110#section-9.2.2 IDEMPOTENT_METHODS = frozenset({"GET", "HEAD", "OPTIONS", "TRACE", "PUT", "DELETE"}) -_RetType = TypeVar("_RetType") +_RetType = TypeVar("_RetType", ClientResponse, ClientWebSocketResponse) _CharsetResolver = Callable[[ClientResponse, bytes], str] From 0f6b71a04f2b61a61572691b8a8677f5ec4d423d Mon Sep 17 00:00:00 2001 From: Sam Bull Date: Sun, 22 Sep 2024 21:49:15 +0100 Subject: [PATCH 10/11] Apply suggestions from code review --- tests/test_client_functional.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_client_functional.py b/tests/test_client_functional.py index c8cfe0e6b62..b8c910914cc 100644 --- a/tests/test_client_functional.py +++ b/tests/test_client_functional.py @@ -3701,7 +3701,7 @@ async def handler(request: web.Request) -> web.Response: async with client.get("/") as resp: with pytest.raises(ClientResponseError, match="404"): # Should not release response when in async with context. - await resp.raise_for_status() + resp.raise_for_status() result = await resp.read() assert result == b"data" @@ -3719,7 +3719,7 @@ async def handler(request: web.Request) -> web.Response: resp = await client.get("/") with pytest.raises(ClientResponseError, match="404"): # No async with, so should release and therefore read() will fail. - await resp.raise_for_status() + resp.raise_for_status() with pytest.raises(aiohttp.ClientConnectionError, match=r"^Connection closed$"): await resp.read() From 2c56bb0ce9ea9f4a7d4bd746cbd85e20e9020076 Mon Sep 17 00:00:00 2001 From: Sam Bull Date: Sun, 22 Sep 2024 22:54:23 +0100 Subject: [PATCH 11/11] Fix --- aiohttp/client.py | 4 ++-- aiohttp/client_ws.py | 14 +++++++++++++- 2 files changed, 15 insertions(+), 3 deletions(-) diff --git a/aiohttp/client.py b/aiohttp/client.py index 39b6b337f0a..89babd47fea 100644 --- a/aiohttp/client.py +++ b/aiohttp/client.py @@ -1275,7 +1275,7 @@ class _BaseRequestContextManager(Coroutine[Any, Any, _RetType], Generic[_RetType __slots__ = ("_coro", "_resp") def __init__(self, coro: Coroutine["asyncio.Future[Any]", None, _RetType]) -> None: - self._coro = coro + self._coro: Coroutine["asyncio.Future[Any]", None, _RetType] = coro def send(self, arg: None) -> "asyncio.Future[Any]": return self._coro.send(arg) @@ -1294,7 +1294,7 @@ def __iter__(self) -> Generator[Any, None, _RetType]: return self.__await__() async def __aenter__(self) -> _RetType: - self._resp = await self._coro + self._resp: _RetType = await self._coro return await self._resp.__aenter__() async def __aexit__( diff --git a/aiohttp/client_ws.py b/aiohttp/client_ws.py index 489cce43992..7e39cc906b0 100644 --- a/aiohttp/client_ws.py +++ b/aiohttp/client_ws.py @@ -3,7 +3,8 @@ import asyncio import dataclasses import sys -from typing import Any, Final, Optional, cast +from types import TracebackType +from typing import Any, Final, Optional, Type, cast from .client_exceptions import ClientError, ServerTimeoutError from .client_reqrep import ClientResponse @@ -395,3 +396,14 @@ async def __anext__(self) -> WSMessage: if msg.type in (WSMsgType.CLOSE, WSMsgType.CLOSING, WSMsgType.CLOSED): raise StopAsyncIteration return msg + + async def __aenter__(self) -> "ClientWebSocketResponse": + return self + + async def __aexit__( + self, + exc_type: Optional[Type[BaseException]], + exc_val: Optional[BaseException], + exc_tb: Optional[TracebackType], + ) -> None: + await self.close()