diff --git a/strawberry/subscriptions/protocols/graphql_transport_ws/handlers.py b/strawberry/subscriptions/protocols/graphql_transport_ws/handlers.py index e9741b9d25..ceec00b4eb 100644 --- a/strawberry/subscriptions/protocols/graphql_transport_ws/handlers.py +++ b/strawberry/subscriptions/protocols/graphql_transport_ws/handlers.py @@ -107,7 +107,7 @@ def on_request_accepted(self) -> None: async def handle_connection_init_timeout(self) -> None: task = asyncio.current_task() - assert task + assert task is not None # for typecheckers try: delay = self.connection_init_wait_timeout.total_seconds() await asyncio.sleep(delay=delay) @@ -281,18 +281,20 @@ async def operation_task(self, operation: Operation) -> None: Operation task top level method. Cleans up and de-registers the operation once it is done. """ - # TODO: Handle errors in this method using self.handle_task_exception() + task = asyncio.current_task() + assert task is not None # for type checkers try: await self.handle_operation(operation) - except BaseException: # pragma: no cover - # cleanup in case of something really unexpected - if operation.id in self.operations: - del self.operations[operation.id] + except asyncio.CancelledError: raise + except Exception as error: + await self.handle_task_exception(error) + # cleanup in case of something really unexpected finally: # add this task to a list to be reaped later - task = asyncio.current_task() - assert task is not None + if operation.id in self.operations: + del self.operations[operation.id] + # TODO: Stop collecting background tasks, not necessary. self.completed_tasks.append(task) async def handle_operation(