File tree 2 files changed +56
-7
lines changed
Filter options
2 files changed +56
-7
lines changed
Original file line number Diff line number Diff line change 2
2
3
3
from __future__ import annotations
4
4
5
- from asyncio import ensure_future , gather , shield , wait_for
5
+ from asyncio import (
6
+ FIRST_EXCEPTION ,
7
+ create_task ,
8
+ ensure_future ,
9
+ gather ,
10
+ shield ,
11
+ wait ,
12
+ wait_for ,
13
+ )
6
14
from contextlib import suppress
7
15
from copy import copy
8
16
from typing import (
@@ -459,12 +467,16 @@ async def get_results() -> dict[str, Any]:
459
467
field = awaitable_fields [0 ]
460
468
results [field ] = await results [field ]
461
469
else :
462
- results .update (
463
- zip (
464
- awaitable_fields ,
465
- await gather (* (results [field ] for field in awaitable_fields )),
466
- )
467
- )
470
+ tasks = {}
471
+ for field in awaitable_fields :
472
+ tasks [create_task (results [field ])] = field # type: ignore[arg-type]
473
+
474
+ done , pending = await wait (tasks , return_when = FIRST_EXCEPTION )
475
+ for task in pending :
476
+ task .cancel ()
477
+
478
+ results .update ((tasks [task ], task .result ()) for task in done )
479
+
468
480
return results
469
481
470
482
return get_results ()
Original file line number Diff line number Diff line change 11
11
GraphQLInt ,
12
12
GraphQLInterfaceType ,
13
13
GraphQLList ,
14
+ GraphQLNonNull ,
14
15
GraphQLObjectType ,
15
16
GraphQLSchema ,
16
17
GraphQLString ,
@@ -193,3 +194,39 @@ async def is_type_of_baz(obj, *_args):
193
194
{"foo" : [{"foo" : "bar" , "foobar" : 1 }, {"foo" : "baz" , "foobaz" : 2 }]},
194
195
None ,
195
196
)
197
+
198
+ @pytest .mark .asyncio
199
+ async def cancel_on_exception ():
200
+ barrier = Barrier (2 )
201
+ completed = asyncio .Event ()
202
+
203
+ async def succeed (* _args ):
204
+ await barrier .wait ()
205
+ completed .set ()
206
+
207
+ async def fail (* _args ):
208
+ raise Exception
209
+
210
+ schema = GraphQLSchema (
211
+ GraphQLObjectType (
212
+ "Query" ,
213
+ {
214
+ "foo" : GraphQLField (GraphQLNonNull (GraphQLBoolean ), resolve = fail ),
215
+ "bar" : GraphQLField (GraphQLBoolean , resolve = succeed ),
216
+ },
217
+ )
218
+ )
219
+
220
+ ast = parse ("{foo, bar}" )
221
+
222
+ awaitable_result = execute (schema , ast )
223
+ assert isinstance (awaitable_result , Awaitable )
224
+ result = await asyncio .wait_for (awaitable_result , 1.0 )
225
+
226
+ assert result .errors
227
+ assert not result .data
228
+
229
+ # Unblock succeed() and check that it does not complete
230
+ await barrier .wait ()
231
+ await asyncio .sleep (0 )
232
+ assert not completed .is_set ()
You can’t perform that action at this time.
0 commit comments