From e5d2a67f9da2d0fe6871fd622d73a10738da218e Mon Sep 17 00:00:00 2001
From: edraft <sven.heidemann@sh-edraft.de>
Date: Tue, 11 Mar 2025 21:40:51 +0100
Subject: [PATCH] Added websocket logging

---
 api/src/api/middleware/websocket.py | 16 ++++++++++++++--
 api/src/core/logger.py              |  4 ++--
 2 files changed, 16 insertions(+), 4 deletions(-)

diff --git a/api/src/api/middleware/websocket.py b/api/src/api/middleware/websocket.py
index e9d66cb..ecd1063 100644
--- a/api/src/api/middleware/websocket.py
+++ b/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
+
+
diff --git a/api/src/core/logger.py b/api/src/core/logger.py
index 1f4a308..0e0c0cc 100644
--- a/api/src/core/logger.py
+++ b/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)