improve

Author: sanderegg

packages/service-library/src/servicelib/fastapi/cancellation_middleware.py

@@ -2,6 +2,7 @@
 import logging
 
 from servicelib.logging_utils import log_context
+from starlette.requests import Request
 from starlette.types import ASGIApp, Receive, Scope, Send
 
 _logger = logging.getLogger(__name__)
@@ -11,10 +12,11 @@ class _TerminateTaskGroupError(Exception):
     pass
 
 
-async def _message_poller(queue, receive):
+async def _message_poller(request: Request, queue: asyncio.Queue, receive: Receive):
     while True:
         message = await receive()
         if message["type"] == "http.disconnect":
+            _logger.info("client disconnected, terminating request to %s!", request.url)
             raise _TerminateTaskGroupError
 
         # Puts the message in the queue
@@ -37,15 +39,21 @@ async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
         # Let's make a shared queue for the request messages
         queue = asyncio.Queue()
 
-        with log_context(_logger, logging.DEBUG, f"cancellable request {scope}"):
+        request = Request(scope)
+
+        with log_context(_logger, logging.DEBUG, f"cancellable request {request.url}"):
             try:
                 async with asyncio.TaskGroup() as tg:
                     handler_task = tg.create_task(
                         _handler(self.app, scope, queue, send)
                     )
-                    poller_task = tg.create_task(_message_poller(queue, receive))
+                    poller_task = tg.create_task(
+                        _message_poller(request, queue, receive)
+                    )
                     response = await handler_task
                     poller_task.cancel()
                     return response
             except* _TerminateTaskGroupError:
-                _logger.info("The client disconnected. The task group was cancelled.")
+                _logger.info(
+                    "The client disconnected. request to %s was cancelled.", request.url
+                )