cancel execution on triggered abort signal despite hanging async resolvers

Replicates graphql/graphql-js@b0b8abe

Author: Cito

src/graphql/execution/execute.py

@@ -3,9 +3,12 @@
 from __future__ import annotations
 
 from asyncio import (
+    FIRST_COMPLETED,
+    CancelledError,
     TimeoutError,  # only needed for Python < 3.11  # noqa: A004
     ensure_future,
     sleep,
+    wait,
 )
 from collections.abc import (
     AsyncGenerator,
@@ -927,6 +930,40 @@ def complete_value(
         )  # pragma: no cover
         raise TypeError(msg)  # pragma: no cover
 
+    async def with_abort_signal(self, awaitable: Awaitable[T]) -> T:
+        """Await a value, but cancel immediately if the abort signal is triggered.
+
+        This wraps awaitables returned by resolvers (and awaitable list items) so
+        that a triggered abort signal interrupts execution *immediately* instead of
+        only at the next field boundary. Without this, a hanging asynchronous
+        resolver would prevent the operation from ever being cancelled.
+
+        If the abort signal fires before the awaitable settles, the underlying
+        awaitable is cancelled and the abort reason is raised (an exception reason
+        is raised as is, any other value is reported as an unexpected error value).
+        """
+        abort_signal = self.abort_signal
+        if abort_signal is None:
+            return await awaitable
+        task = ensure_future(awaitable)
+        if not abort_signal.aborted:
+            abort = ensure_future(abort_signal.wait())
+            try:
+                await wait({task, abort}, return_when=FIRST_COMPLETED)
+            finally:
+                if not abort.done():
+                    abort.cancel()
+            if task.done():
+                return task.result()
+        task.cancel()
+        with suppress(CancelledError):
+            await task
+        reason = abort_signal.reason
+        if isinstance(reason, Exception):
+            raise reason
+        msg = f"Unexpected error value: {inspect(reason)}"
+        raise TypeError(msg)
+
     async def complete_awaitable_value(
         self,
         return_type: GraphQLOutputType,
@@ -939,7 +976,7 @@ async def complete_awaitable_value(
     ) -> GraphQLWrappedResult[Any]:
         """Complete an awaitable value."""
         try:
-            resolved = await result
+            resolved = await self.with_abort_signal(result)
             completed = self.complete_value(
                 return_type,
                 field_details_list,
@@ -1363,7 +1400,7 @@ async def complete_awaitable_list_item_value(
     ) -> Any:
         """Complete an awaitable list item value."""
         try:
-            resolved = await item
+            resolved = await self.with_abort_signal(item)
             completed = self.complete_value(
                 item_type,
                 field_details_list,
@@ -1527,14 +1564,6 @@ def complete_object_value(
         defer_map: RefMap[DeferUsage, DeferredFragmentRecord] | None,
     ) -> AwaitableOrValue[GraphQLWrappedResult[dict[str, Any]]]:
         """Complete an Object value by executing all sub-selections."""
-        abort_signal = self.abort_signal
-        if abort_signal is not None and abort_signal.aborted:
-            raise located_error(
-                abort_signal.reason,
-                to_nodes(field_details_list),
-                path.as_list(),
-            )
-
         # If there is an `is_type_of()` predicate function, call it with the current
         # result. If `is_type_of()` returns False, then raise an error rather than
         # continuing execution.
@@ -2845,7 +2874,7 @@ def execute_subscription(
 
             async def await_result() -> AsyncIterable[Any]:
                 try:
-                    return assert_event_stream(await result)
+                    return assert_event_stream(await context.with_abort_signal(result))
                 except Exception as error:
                     raise located_error(error, field_nodes, path.as_list()) from error
 

tests/execution/test_abort_signal.py

@@ -2,13 +2,13 @@
 
 from __future__ import annotations
 
-from asyncio import ensure_future, sleep
+from asyncio import Event, Future, ensure_future, sleep
 from collections.abc import Awaitable
 
 import pytest
 
 from graphql import build_schema
-from graphql.execution import execute
+from graphql.execution import execute, subscribe
 from graphql.language import parse
 from graphql.pyutils import AbortController, AbortError
 
@@ -19,7 +19,7 @@
     """
     type Todo {
       id: ID
-      text: String
+      items: [String]
       author: User
     }
 
@@ -30,12 +30,17 @@
 
     type Query {
       todo: Todo
+      nonNullableTodo: Todo!
     }
 
     type Mutation {
       foo: String
       bar: String
     }
+
+    type Subscription {
+      foo: String
+    }
     """
 )
 
@@ -97,7 +102,6 @@ async def stops_the_execution_when_aborted_during_object_field_completion():
         async def todo(_info):
             return {
                 "id": "1",
-                "text": "Hello, World!",
                 "author": must_not_be_called,
             }
 
@@ -190,7 +194,6 @@ async def stops_the_execution_when_aborted_during_completion_with_custom_error()
         async def todo(_info):
             return {
                 "id": "1",
-                "text": "Hello, World!",
                 "author": must_not_be_called,
             }
 
@@ -238,7 +241,6 @@ async def stops_the_execution_when_aborted_during_completion_with_custom_string(
         async def todo(_info):
             return {
                 "id": "1",
-                "text": "Hello, World!",
                 "author": must_not_be_called,
             }
 
@@ -290,7 +292,6 @@ async def author(_info):
             root_value={
                 "todo": {
                     "id": "1",
-                    "text": "Hello, World!",
                     "author": author,
                 }
             },
@@ -312,6 +313,146 @@ async def author(_info):
             ],
         )
 
+    async def stops_the_execution_when_aborted_despite_a_hanging_resolver():
+        abort_controller = AbortController()
+        document = parse(
+            """
+      query {
+        todo {
+          id
+          author {
+            id
+          }
+        }
+      }
+    """
+        )
+
+        started = Event()
+
+        async def todo(_info):
+            started.set()
+            await Future()  # will never resolve
+
+        awaitable_result = execute(
+            schema,
+            document,
+            abort_signal=abort_controller.signal,
+            root_value={"todo": todo},
+        )
+        assert isinstance(awaitable_result, Awaitable)
+
+        # Abort only once the resolver is actually in flight, so that cancellation
+        # must interrupt the hanging resolver instead of being caught up front.
+        task = ensure_future(awaitable_result)
+        await started.wait()
+        abort_controller.abort()
+
+        result = await task
+
+        assert result.errors is not None
+        assert isinstance(result.errors[0].original_error, AbortError)
+        assert result == (
+            {"todo": None},
+            [
+                {
+                    "message": "This operation was aborted",
+                    "locations": [(3, 9)],
+                    "path": ["todo"],
+                }
+            ],
+        )
+
+    async def stops_the_execution_when_aborted_despite_a_hanging_item():
+        abort_controller = AbortController()
+        document = parse(
+            """
+      query {
+        todo {
+          id
+          items
+        }
+      }
+    """
+        )
+
+        def todo(_info):
+            return {
+                "id": "1",
+                "items": [Future()],  # will never resolve
+            }
+
+        awaitable_result = execute(
+            schema,
+            document,
+            abort_signal=abort_controller.signal,
+            root_value={"todo": todo},
+        )
+        assert isinstance(awaitable_result, Awaitable)
+
+        abort_controller.abort()
+
+        result = await awaitable_result
+
+        assert result.errors is not None
+        assert isinstance(result.errors[0].original_error, AbortError)
+        assert result == (
+            {"todo": {"id": "1", "items": [None]}},
+            [
+                {
+                    "message": "This operation was aborted",
+                    "locations": [(5, 11)],
+                    "path": ["todo", "items", 0],
+                }
+            ],
+        )
+
+    async def stops_the_execution_when_aborted_with_proper_null_bubbling():
+        abort_controller = AbortController()
+        document = parse(
+            """
+      query {
+        nonNullableTodo {
+          id
+          author {
+            id
+          }
+        }
+      }
+    """
+        )
+
+        async def non_nullable_todo(_info):
+            return {
+                "id": "1",
+                "author": must_not_be_called,
+            }
+
+        awaitable_result = execute(
+            schema,
+            document,
+            abort_signal=abort_controller.signal,
+            root_value={"nonNullableTodo": non_nullable_todo},
+        )
+        assert isinstance(awaitable_result, Awaitable)
+
+        abort_controller.abort()
+
+        result = await awaitable_result
+
+        assert result.errors is not None
+        assert isinstance(result.errors[0].original_error, AbortError)
+        assert result == (
+            None,
+            [
+                {
+                    "message": "This operation was aborted",
+                    "locations": [(3, 9)],
+                    "path": ["nonNullableTodo"],
+                }
+            ],
+        )
+
     async def stops_the_execution_when_aborted_mid_mutation():
         abort_controller = AbortController()
         document = parse(
@@ -334,9 +475,16 @@ async def foo(_info):
         )
         assert isinstance(awaitable_result, Awaitable)
 
+        # Let the first field resolve before aborting, so that the abort is only
+        # observed when serially moving on to the second field (mirrors the
+        # ``resolveOnNextTick`` calls in the GraphQL.js test).
+        task = ensure_future(awaitable_result)
+        for _ in range(3):
+            await sleep(0)
+
         abort_controller.abort()
 
-        result = await awaitable_result
+        result = await task
 
         assert result == (
             {"foo": "baz", "bar": None},
@@ -373,3 +521,39 @@ async def stops_the_execution_when_aborted_pre_execute():
         )
 
         assert result == (None, [{"message": "This operation was aborted"}])
+
+    async def stops_the_execution_when_aborted_during_subscription():
+        abort_controller = AbortController()
+        document = parse(
+            """
+      subscription {
+        foo
+      }
+    """
+        )
+
+        def foo(_info):
+            return Future()  # will never resolve
+
+        awaitable_result = subscribe(
+            schema,
+            document,
+            abort_signal=abort_controller.signal,
+            root_value={"foo": foo},
+        )
+        assert isinstance(awaitable_result, Awaitable)
+
+        abort_controller.abort()
+
+        result = await awaitable_result
+
+        assert result == (
+            None,
+            [
+                {
+                    "message": "This operation was aborted",
+                    "locations": [(3, 9)],
+                    "path": ["foo"],
+                }
+            ],
+        )