Cancel remaining list items on exceptions

Author: Cito

src/graphql/execution/execute.py

@@ -466,10 +466,10 @@ async def get_results() -> dict[str, Any]:
                 field = awaitable_fields[0]
                 results[field] = await results[field]
             else:
-                tasks = {
-                    create_task(results[field]): field  # type: ignore[arg-type]
+                tasks = [
+                    create_task(results[field])  # type: ignore[arg-type]
                     for field in awaitable_fields
-                }
+                ]
 
                 try:
                     awaited_results = await gather(*tasks)
@@ -1014,12 +1014,21 @@ async def get_completed_results() -> list[Any]:
                 index = awaitable_indices[0]
                 completed_results[index] = await completed_results[index]
             else:
-                for index, sub_result in zip(
-                    awaitable_indices,
-                    await gather(
-                        *(completed_results[index] for index in awaitable_indices)
-                    ),
-                ):
+                tasks = [
+                    create_task(completed_results[index]) for index in awaitable_indices
+                ]
+
+                try:
+                    awaited_results = await gather(*tasks)
+                except Exception:
+                    # Cancel unfinished tasks before raising the exception
+                    for task in tasks:
+                        if not task.done():
+                            task.cancel()
+                    await gather(*tasks, return_exceptions=True)
+                    raise
+
+                for index, sub_result in zip(awaitable_indices, awaited_results):
                     completed_results[index] = sub_result
             return completed_results
 

tests/execution/test_parallel.py

@@ -76,7 +76,7 @@ async def resolve(*_args):
         # raises TimeoutError if not parallel
         awaitable_result = execute(schema, ast)
         assert isinstance(awaitable_result, Awaitable)
-        result = await asyncio.wait_for(awaitable_result, 1.0)
+        result = await asyncio.wait_for(awaitable_result, 1)
 
         assert result == ({"foo": True, "bar": True}, None)
 
@@ -125,7 +125,7 @@ async def resolve_list(*args):
         # raises TimeoutError if not parallel
         awaitable_result = execute(schema, ast)
         assert isinstance(awaitable_result, Awaitable)
-        result = await asyncio.wait_for(awaitable_result, 1.0)
+        result = await asyncio.wait_for(awaitable_result, 1)
 
         assert result == ({"foo": [True, True]}, None)
 
@@ -188,15 +188,15 @@ async def is_type_of_baz(obj, *_args):
         # raises TimeoutError if not parallel
         awaitable_result = execute(schema, ast)
         assert isinstance(awaitable_result, Awaitable)
-        result = await asyncio.wait_for(awaitable_result, 1.0)
+        result = await asyncio.wait_for(awaitable_result, 1)
 
         assert result == (
             {"foo": [{"foo": "bar", "foobar": 1}, {"foo": "baz", "foobaz": 2}]},
             None,
         )
 
     @pytest.mark.asyncio
-    async def cancel_on_exception():
+    async def cancel_selection_sets_on_exception():
         barrier = Barrier(2)
         completed = False
 
@@ -222,13 +222,61 @@ async def fail(*_args):
 
         awaitable_result = execute(schema, ast)
         assert isinstance(awaitable_result, Awaitable)
-        result = await asyncio.wait_for(awaitable_result, 1.0)
+        result = await asyncio.wait_for(awaitable_result, 1)
 
         assert result == (
             None,
             [{"message": "Oops", "locations": [(1, 2)], "path": ["foo"]}],
         )
 
+        assert not completed
+
+        # Unblock succeed() and check that it does not complete
+        await barrier.wait()
+        await asyncio.sleep(0)
+        assert not completed
+
+    @pytest.mark.asyncio
+    async def cancel_lists_on_exception():
+        barrier = Barrier(2)
+        completed = False
+
+        async def succeed(*_args):
+            nonlocal completed
+            await barrier.wait()
+            completed = True  # pragma: no cover
+
+        async def fail(*_args):
+            raise RuntimeError("Oops")
+
+        async def resolve_list(*args):
+            return [fail(*args), succeed(*args)]
+
+        schema = GraphQLSchema(
+            GraphQLObjectType(
+                "Query",
+                {
+                    "foo": GraphQLField(
+                        GraphQLList(GraphQLNonNull(GraphQLBoolean)),
+                        resolve=resolve_list,
+                    )
+                },
+            )
+        )
+
+        ast = parse("{foo}")
+
+        awaitable_result = execute(schema, ast)
+        assert isinstance(awaitable_result, Awaitable)
+        result = await asyncio.wait_for(awaitable_result, 1)
+
+        assert result == (
+            {"foo": None},
+            [{"message": "Oops", "locations": [(1, 2)], "path": ["foo", 0]}],
+        )
+
+        assert not completed
+
         # Unblock succeed() and check that it does not complete
         await barrier.wait()
         await asyncio.sleep(0)