Merge pull request #113 from tartiflette/ISSUE-112

ISSUE-112 - Use context_factory as an async ctx manager like for a normal query
tartiflette · Sep 24, 2020 · 1f9a27e · 1f9a27e
2 parents 1a6e05f + d3902a4
commit 1f9a27e
Showing 1 changed file with 5 additions and 4 deletions.
diff --git a/tartiflette_aiohttp/_subscription_ws_handler.py b/tartiflette_aiohttp/_subscription_ws_handler.py
@@ -262,7 +262,8 @@ async def _handle_request(self) -> None:
 
     async def __call__(self, request: "Request") -> "WebSocketResponse":
         self._socket = web.WebSocketResponse(protocols=(WS_PROTOCOL,))
-        self._context = await self._context_factory(request)
-        await self._socket.prepare(request)
-        await shield(self._handle_request())
-        return self._socket
+        async with self._context_factory(request) as ctx:
+            self._context = ctx
+            await self._socket.prepare(request)
+            await shield(self._handle_request())
+            return self._socket