Make AIOHTTPTransport.connect() reentrant for shared client instances

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,16 +141,23 @@ 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.
 
         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")

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())