Added websocket logging

api/src/api/middleware/websocket.py

@@ -1,3 +1,5 @@
+from uuid import uuid4
+
 from ariadne.asgi.handlers import GraphQLTransportWSHandler
 from starlette.datastructures import MutableHeaders
 from starlette.websockets import WebSocket
@@ -5,16 +7,19 @@ from starlette.websockets import WebSocket
 from api.middleware.request import set_request
 from core.logger import APILogger
 
-logger = APILogger("WS")
+logger = APILogger("api.ws")
 
 
 class AuthenticatedGraphQLTransportWSHandler(GraphQLTransportWSHandler):
 
     def __init__(self, *args, **kwargs):
-        super().__init__(*args, on_connect=self.on_connect, **kwargs)
+        super().__init__(*args, on_connect=self.on_connect, on_disconnect=self.on_disconnect, **kwargs)
 
     @staticmethod
     async def on_connect(ws: WebSocket, message: dict):
+        ws.state.request_id = uuid4()
+        logger.info(f"WebSocket connection {ws.state.request_id}")
+
         if "Authorization" not in message:
             return True
 
@@ -24,3 +29,10 @@ class AuthenticatedGraphQLTransportWSHandler(GraphQLTransportWSHandler):
 
         set_request(ws)
         return True
+
+    @staticmethod
+    async def on_disconnect(ws: WebSocket):
+        logger.debug(f"WebSocket connection {ws.state.request_id} closed")
+        return True
+
+

api/src/core/logger.py

@@ -72,8 +72,8 @@ class Logger:
         if request is not None:
             structured_message["request"] = {
                 "url": str(request.url),
-                "method": request.method,
-                "data": asyncio.create_task(request.body()),
+                "method": request.method if request.scope == "http" else "ws",
+                "data": asyncio.create_task(request.body()) if request.scope == "http" else None,
             }
         return str(structured_message)