Repository: mirumee/ariadne

Issue:
on_disconnect websocket handler is not called for all possible disconnects
Currently `on_disonnect` callback is only called while handling `GQL_CONNECTION_TERMINATE` message, but the websocket can be disconnected in other places:
- exception in `handle_websocket_connection_init_message`
- on WebSocketDisconnect in `websocket_server`

Calling on_disconnect in those cases would be very useful to properly handle websocket related cleanups


Discussion:
This is a little bit painful to us too, here is our workaround:

```
class GraphQLApp(GraphQL):
    async def websocket_server(self, websocket: WebSocket) -> None:
        await super().websocket_server(websocket)
        await db.aclose()
```

with a huge comment block why it's necessary here, and why we can't use `on_disconnect`.
@kissgyorgy it feels wrong to open/close service conenctions in websocket hooks. Every ASGI lib I've seen so far created connections pool on startup and picked connection from this pool on demand. Connections were closed only on app's shutdown.
We have a multi-tenant architecture, we can't have a pool, because the database is determined on connection from the client, so every client has a different connection which has to be closed when the websocket connection is closed, otherwise we would have hundreds of connections lingering unnecessarily. (Terrible legacy architecture, but we inherited it, not easy to change.)

Resolving patch:
diff --git a/ariadne/asgi.py b/ariadne/asgi.py
--- a/ariadne/asgi.py
+++ b/ariadne/asgi.py
@@ -1,5 +1,6 @@
 import asyncio
 import json
+from dataclasses import dataclass
 from inspect import isawaitable
 from typing import (
     Any,
@@ -38,7 +39,6 @@
     ValidationRules,
 )
 
-
 GQL_CONNECTION_INIT = "connection_init"  # Client -> Server
 GQL_CONNECTION_ACK = "connection_ack"  # Server -> Client
 GQL_CONNECTION_ERROR = "connection_error"  # Server -> Client
@@ -73,8 +73,17 @@ def __init__(self, payload: Union[dict, str] = None):
 ]
 
 
-OnConnect = Callable[[WebSocket, Any], None]
-OnDisconnect = Callable[[WebSocket], None]
+@dataclass
+class Operation:
+    id: str
+    name: Optional[str]
+    generator: AsyncGenerator
+
+
+OnConnect = Callable[[WebSocket, Any], Any]
+OnDisconnect = Callable[[WebSocket], Any]
+OnOperation = Callable[[WebSocket, Operation], Any]
+OnComplete = Callable[[WebSocket, Operation], Any]
 
 
 class GraphQL:
@@ -86,6 +95,8 @@ def __init__(
         root_value: Optional[RootValue] = None,
         on_connect: Optional[OnConnect] = None,
         on_disconnect: Optional[OnDisconnect] = None,
+        on_operation: Optional[OnOperation] = None,
+        on_complete: Optional[OnComplete] = None,
         validation_rules: Optional[ValidationRules] = None,
         debug: bool = False,
         introspection: bool = True,
@@ -99,6 +110,8 @@ def __init__(
         self.root_value = root_value
         self.on_connect = on_connect
         self.on_disconnect = on_disconnect
+        self.on_operation = on_operation
+        self.on_complete = on_complete
         self.validation_rules = validation_rules
         self.debug = debug
         self.introspection = introspection
@@ -268,7 +281,7 @@ def handle_not_allowed_method(self, request: Request):
         return Response(status_code=405, headers=allow_header)
 
     async def websocket_server(self, websocket: WebSocket) -> None:
-        subscriptions: Dict[str, AsyncGenerator] = {}
+        operations: Dict[str, Operation] = {}
         await websocket.accept("graphql-ws")
         try:
             while WebSocketState.DISCONNECTED not in (
@@ -276,18 +289,29 @@ async def websocket_server(self, websocket: WebSocket) -> None:
                 websocket.application_state,
             ):
                 message = await websocket.receive_json()
-                await self.handle_websocket_message(message, websocket, subscriptions)
+                await self.handle_websocket_message(message, websocket, operations)
         except WebSocketDisconnect:
             pass
         finally:
-            for subscription in subscriptions.values():
-                await subscription.aclose()
+            for operation_id, operation in list(operations.items()):
+                await self.stop_websocket_operation(websocket, operation)
+                del operations[operation_id]
+
+            try:
+                if self.on_disconnect:
+                    result = self.on_disconnect(websocket)
+                    if result and isawaitable(result):
+                        await result
+            except Exception as error:
+                if not isinstance(error, GraphQLError):
+                    error = GraphQLError(str(error), original_error=error)
+                log_error(error, self.logger)
 
     async def handle_websocket_message(
         self,
         message: dict,
         websocket: WebSocket,
-        subscriptions: Dict[str, AsyncGenerator],
+        operations: Dict[str, Operation],
     ):
         operation_id = cast(str, message.get("id"))
         message_type = cast(str, message.get("type"))
@@ -302,13 +326,13 @@ async def handle_websocket_message(
                 websocket,
             )
         elif message_type == GQL_START:
-            await self.start_websocket_subscription(
-                message.get("payload"), operation_id, websocket, subscriptions
+            await self.start_websocket_operation(
+                message.get("payload"), operation_id, websocket, operations
             )
         elif message_type == GQL_STOP:
-            if operation_id in subscriptions:
-                await subscriptions[operation_id].aclose()
-                del subscriptions[operation_id]
+            if operation_id in operations:
+                await self.stop_websocket_operation(websocket, operations[operation_id])
+                del operations[operation_id]
 
     async def handle_websocket_connection_init_message(
         self,
@@ -342,16 +366,6 @@ async def handle_websocket_connection_terminate_message(
     ):
         await websocket.close()
 
-        try:
-            if self.on_disconnect:
-                result = self.on_disconnect(websocket)
-                if result and isawaitable(result):
-                    await result
-        except Exception as error:
-            if not isinstance(error, GraphQLError):
-                error = GraphQLError(str(error), original_error=error)
-            log_error(error, self.logger)
-
     async def keep_websocket_alive(self, websocket: WebSocket):
         if not self.keepalive:
             return
@@ -362,12 +376,12 @@ async def keep_websocket_alive(self, websocket: WebSocket):
                 return
             await asyncio.sleep(self.keepalive)
 
-    async def start_websocket_subscription(
+    async def start_websocket_operation(
         self,
         data: Any,
         operation_id: str,
         websocket: WebSocket,
-        subscriptions: Dict[str, AsyncGenerator],
+        operations: Dict[str, Operation],
     ):
         context_value = await self.get_context_for_request(websocket)
 
@@ -389,11 +403,41 @@ async def start_websocket_subscription(
             )
         else:
             results = cast(AsyncGenerator, results)
-            subscriptions[operation_id] = results
+            operations[operation_id] = Operation(
+                id=operation_id,
+                name=data.get("operationName"),
+                generator=results,
+            )
+
+            if self.on_operation:
+                try:
+                    result = self.on_operation(websocket, operations[operation_id])
+                    if result and isawaitable(result):
+                        await result
+                except Exception as error:
+                    if not isinstance(error, GraphQLError):
+                        error = GraphQLError(str(error), original_error=error)
+                    log_error(error, self.logger)
+
             asyncio.ensure_future(
                 self.observe_async_results(results, operation_id, websocket)
             )
 
+    async def stop_websocket_operation(
+        self, websocket: WebSocket, operation: Operation
+    ):
+        if self.on_complete:
+            try:
+                result = self.on_complete(websocket, operation)
+                if result and isawaitable(result):
+                    await result
+            except Exception as error:
+                if not isinstance(error, GraphQLError):
+                    error = GraphQLError(str(error), original_error=error)
+                log_error(error, self.logger)
+
+        await operation.generator.aclose()
+
     async def observe_async_results(
         self, results: AsyncGenerator, operation_id: str, websocket: WebSocket
     ) -> None:


Test patch:
diff --git a/tests/asgi/snapshots/snap_test_query_execution.py b/tests/asgi/snapshots/snap_test_query_execution.py
--- a/tests/asgi/snapshots/snap_test_query_execution.py
+++ b/tests/asgi/snapshots/snap_test_query_execution.py
@@ -93,7 +93,7 @@
 
 snapshots['test_query_is_executed_for_multipart_request_with_large_file_with_tracing 1'] = {
     'data': {
-        'upload': 'UploadFile' 
+        'upload': 'UploadFile'
     }
 }
 
diff --git a/tests/asgi/test_websockets.py b/tests/asgi/test_websockets.py
--- a/tests/asgi/test_websockets.py
+++ b/tests/asgi/test_websockets.py
@@ -204,7 +204,217 @@ def on_connect(websocket, payload):
         assert response["payload"] == {"msg": "Token required", "code": "auth_error"}
 
 
-def test_custom_websocket_on_disconnect_is_called(schema):
+def test_custom_websocket_on_operation_is_called(schema):
+    def on_operation(websocket, operation):
+        assert operation.name == "TestOp"
+        websocket.scope["on_operation"] = True
+
+    app = GraphQL(schema, on_operation=on_operation)
+    client = TestClient(app)
+
+    with client.websocket_connect("/", "graphql-ws") as ws:
+        ws.send_json({"type": GQL_CONNECTION_INIT})
+        response = ws.receive_json()
+        assert response["type"] == GQL_CONNECTION_ACK
+        ws.send_json(
+            {
+                "type": GQL_START,
+                "id": "test1",
+                "payload": {
+                    "operationName": "TestOp",
+                    "query": "subscription TestOp { ping }",
+                },
+            }
+        )
+        response = ws.receive_json()
+        assert response["type"] == GQL_DATA
+        assert response["id"] == "test1"
+        assert response["payload"]["data"] == {"ping": "pong"}
+        ws.send_json({"type": GQL_STOP})
+        response = ws.receive_json()
+        assert response["type"] == GQL_DATA
+        assert response["id"] == "test1"
+        assert response["payload"]["data"] == {"ping": "ping"}
+        ws.send_json({"type": GQL_CONNECTION_TERMINATE})
+        assert ws.scope["on_operation"] is True
+
+
+def test_custom_websocket_on_operation_is_awaited_if_its_async(schema):
+    async def on_operation(websocket, operation):
+        assert operation.name == "TestOp"
+        websocket.scope["on_operation"] = True
+
+    app = GraphQL(schema, on_operation=on_operation)
+    client = TestClient(app)
+
+    with client.websocket_connect("/", "graphql-ws") as ws:
+        ws.send_json({"type": GQL_CONNECTION_INIT})
+        response = ws.receive_json()
+        assert response["type"] == GQL_CONNECTION_ACK
+        ws.send_json(
+            {
+                "type": GQL_START,
+                "id": "test1",
+                "payload": {
+                    "operationName": "TestOp",
+                    "query": "subscription TestOp { ping }",
+                },
+            }
+        )
+        response = ws.receive_json()
+        assert response["type"] == GQL_DATA
+        assert response["id"] == "test1"
+        assert response["payload"]["data"] == {"ping": "pong"}
+        ws.send_json({"type": GQL_STOP})
+        response = ws.receive_json()
+        assert response["type"] == GQL_DATA
+        assert response["id"] == "test1"
+        assert response["payload"]["data"] == {"ping": "ping"}
+        ws.send_json({"type": GQL_CONNECTION_TERMINATE})
+        assert ws.scope["on_operation"] is True
+
+
+def test_error_in_custom_websocket_on_operation_is_handled(schema):
+    async def on_operation(websocket, operation):
+        raise ValueError("Oh No!")
+
+    app = GraphQL(schema, on_operation=on_operation)
+    client = TestClient(app)
+
+    with client.websocket_connect("/", "graphql-ws") as ws:
+        ws.send_json({"type": GQL_CONNECTION_INIT})
+        response = ws.receive_json()
+        assert response["type"] == GQL_CONNECTION_ACK
+        ws.send_json(
+            {
+                "type": GQL_START,
+                "id": "test1",
+                "payload": {
+                    "operationName": "TestOp",
+                    "query": "subscription TestOp { ping }",
+                },
+            }
+        )
+        response = ws.receive_json()
+        assert response["type"] == GQL_DATA
+        assert response["id"] == "test1"
+        assert response["payload"]["data"] == {"ping": "pong"}
+        ws.send_json({"type": GQL_STOP})
+        response = ws.receive_json()
+        assert response["type"] == GQL_DATA
+        assert response["id"] == "test1"
+        assert response["payload"]["data"] == {"ping": "ping"}
+        ws.send_json({"type": GQL_CONNECTION_TERMINATE})
+
+
+def test_custom_websocket_on_complete_is_called(schema):
+    def on_complete(websocket, operation):
+        assert operation.name == "TestOp"
+        websocket.scope["on_complete"] = True
+
+    app = GraphQL(schema, on_complete=on_complete)
+    client = TestClient(app)
+
+    with client.websocket_connect("/", "graphql-ws") as ws:
+        ws.send_json({"type": GQL_CONNECTION_INIT})
+        response = ws.receive_json()
+        assert response["type"] == GQL_CONNECTION_ACK
+        ws.send_json(
+            {
+                "type": GQL_START,
+                "id": "test1",
+                "payload": {
+                    "operationName": "TestOp",
+                    "query": "subscription TestOp { ping }",
+                },
+            }
+        )
+        response = ws.receive_json()
+        assert response["type"] == GQL_DATA
+        assert response["id"] == "test1"
+        assert response["payload"]["data"] == {"ping": "pong"}
+        ws.send_json({"type": GQL_STOP})
+        response = ws.receive_json()
+        assert response["type"] == GQL_DATA
+        assert response["id"] == "test1"
+        assert response["payload"]["data"] == {"ping": "ping"}
+        ws.send_json({"type": GQL_CONNECTION_TERMINATE})
+        assert "on_complete" not in ws.scope
+
+    assert ws.scope["on_complete"] is True
+
+
+def test_custom_websocket_on_complete_is_awaited_if_its_async(schema):
+    async def on_complete(websocket, operation):
+        assert operation.name == "TestOp"
+        websocket.scope["on_complete"] = True
+
+    app = GraphQL(schema, on_complete=on_complete)
+    client = TestClient(app)
+
+    with client.websocket_connect("/", "graphql-ws") as ws:
+        ws.send_json({"type": GQL_CONNECTION_INIT})
+        response = ws.receive_json()
+        assert response["type"] == GQL_CONNECTION_ACK
+        ws.send_json(
+            {
+                "type": GQL_START,
+                "id": "test1",
+                "payload": {
+                    "operationName": "TestOp",
+                    "query": "subscription TestOp { ping }",
+                },
+            }
+        )
+        response = ws.receive_json()
+        assert response["type"] == GQL_DATA
+        assert response["id"] == "test1"
+        assert response["payload"]["data"] == {"ping": "pong"}
+        ws.send_json({"type": GQL_STOP})
+        response = ws.receive_json()
+        assert response["type"] == GQL_DATA
+        assert response["id"] == "test1"
+        assert response["payload"]["data"] == {"ping": "ping"}
+        ws.send_json({"type": GQL_CONNECTION_TERMINATE})
+        assert "on_complete" not in ws.scope
+
+    assert ws.scope["on_complete"] is True
+
+
+def test_error_in_custom_websocket_on_complete_is_handled(schema):
+    async def on_complete(websocket, operation):
+        raise ValueError("Oh No!")
+
+    app = GraphQL(schema, on_complete=on_complete)
+    client = TestClient(app)
+
+    with client.websocket_connect("/", "graphql-ws") as ws:
+        ws.send_json({"type": GQL_CONNECTION_INIT})
+        response = ws.receive_json()
+        assert response["type"] == GQL_CONNECTION_ACK
+        ws.send_json(
+            {
+                "type": GQL_START,
+                "id": "test1",
+                "payload": {
+                    "operationName": "TestOp",
+                    "query": "subscription TestOp { ping }",
+                },
+            }
+        )
+        response = ws.receive_json()
+        assert response["type"] == GQL_DATA
+        assert response["id"] == "test1"
+        assert response["payload"]["data"] == {"ping": "pong"}
+        ws.send_json({"type": GQL_STOP})
+        response = ws.receive_json()
+        assert response["type"] == GQL_DATA
+        assert response["id"] == "test1"
+        assert response["payload"]["data"] == {"ping": "ping"}
+        ws.send_json({"type": GQL_CONNECTION_TERMINATE})
+
+
+def test_custom_websocket_on_disconnect_is_called_on_terminate_message(schema):
     def on_disconnect(websocket):
         websocket.scope["on_disconnect"] = True
 
@@ -221,6 +431,22 @@ def on_disconnect(websocket):
     assert ws.scope["on_disconnect"] is True
 
 
+def test_custom_websocket_on_disconnect_is_called_on_connection_close(schema):
+    def on_disconnect(websocket):
+        websocket.scope["on_disconnect"] = True
+
+    app = GraphQL(schema, on_disconnect=on_disconnect)
+    client = TestClient(app)
+
+    with client.websocket_connect("/", "graphql-ws") as ws:
+        ws.send_json({"type": GQL_CONNECTION_INIT})
+        response = ws.receive_json()
+        assert response["type"] == GQL_CONNECTION_ACK
+        assert "on_disconnect" not in ws.scope
+
+    assert ws.scope["on_disconnect"] is True
+
+
 def test_custom_websocket_on_disconnect_is_awaited_if_its_async(schema):
     async def on_disconnect(websocket):
         websocket.scope["on_disconnect"] = True
diff --git a/tests/conftest.py b/tests/conftest.py
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -85,6 +85,8 @@ def mutations():
 
 async def ping_generator(*_):
     yield {"ping": "pong"}
+    yield {"ping": "ping"}
+    yield {"ping": "pong"}
 
 
 async def error_generator(*_):