diff --git a/gql/transport/aiohttp.py b/gql/transport/aiohttp.py index e7eff55f..f7f48577 100644 --- a/gql/transport/aiohttp.py +++ b/gql/transport/aiohttp.py @@ -98,6 +98,7 @@ def __init__( self.ssl_close_timeout: Optional[Union[int, float]] = ssl_close_timeout self.client_session_args = client_session_args self.session: Optional[aiohttp.ClientSession] = None + self._connect_count: int = 0 self.response_headers: Optional[CIMultiDictProxy[str]] self.json_serialize: Callable = json_serialize self.json_deserialize: Callable = json_deserialize @@ -110,6 +111,10 @@ async def connect(self) -> None: to create the session. Should be cleaned with a call to the close coroutine. + + This method is reentrant: if the session already exists, the existing + session is reused and a reference count is incremented. Each connect() + call must be paired with a close() call. """ if self.session is None: @@ -136,8 +141,7 @@ async def connect(self) -> None: self.session = aiohttp.ClientSession(**client_session_args) - else: - raise TransportAlreadyConnected("Transport is already connected") + self._connect_count += 1 async def close(self) -> None: """Coroutine which will close the aiohttp session. @@ -145,7 +149,15 @@ async def close(self) -> None: Don't call this coroutine directly on the transport, instead use :code:`async with` on the client and this coroutine will be executed when you exit the async context manager. + + The session is only actually closed when every connect() call has been + balanced by a close() call. """ + self._connect_count = max(0, self._connect_count - 1) + + if self._connect_count > 0: + return + if self.session is not None: log.debug("Closing transport") diff --git a/tests/test_aiohttp.py b/tests/test_aiohttp.py index 00bd8a0f..f80d9865 100644 --- a/tests/test_aiohttp.py +++ b/tests/test_aiohttp.py @@ -1,3 +1,4 @@ +import asyncio import io import json import os @@ -391,7 +392,7 @@ async def handler(request): @pytest.mark.asyncio -async def test_aiohttp_cannot_connect_twice(aiohttp_server): +async def test_aiohttp_connect_twice_is_reentrant(aiohttp_server): from aiohttp import web from gql.transport.aiohttp import AIOHTTPTransport @@ -408,9 +409,12 @@ async def handler(request): transport = AIOHTTPTransport(url=url, timeout=10) async with Client(transport=transport) as session: - - with pytest.raises(TransportAlreadyConnected): - await session.transport.connect() + # Second connect should succeed (reentrant) and reuse the session + original_session = session.transport.session + await session.transport.connect() + assert session.transport.session is original_session + # Balance the extra connect with a close + await session.transport.close() @pytest.mark.asyncio @@ -1918,3 +1922,38 @@ async def handler(request): await session.execute("qmlsdkfj") assert "request should be a GraphQLRequest object" in str(exc_info.value) + + +@pytest.mark.asyncio +async def test_aiohttp_reentrant_connect(aiohttp_server): + """A single Client/transport instance should support concurrent sessions. + + When a Client is cached (e.g. via lru_cache) and used from multiple + concurrent async-with blocks, the second caller should not get + TransportAlreadyConnected. + """ + from aiohttp import web + + from gql.transport.aiohttp import AIOHTTPTransport + + async def handler(request): + return web.Response(text=query1_server_answer, content_type="application/json") + + app = web.Application() + app.router.add_route("POST", "/", handler) + server = await aiohttp_server(app) + + url = server.make_url("/") + + transport = AIOHTTPTransport(url=url, timeout=10) + client = Client(transport=transport) + + results = [] + + async def use_client(): + async with client as session: + query = gql(query1_str) + result = await session.execute(query) + results.append(result) + + await asyncio.gather(use_client(), use_client())