From 0a107ad4482c4cdc2d4ea8494099c061628f59f4 Mon Sep 17 00:00:00 2001 From: fan wu <336798967+FanWu-ai@users.noreply.github.com> Date: Sun, 11 Oct 2026 10:08:33 -0700 Subject: [PATCH] Preserve DataLoader batch completion after a request is cancelled AI assistance: OpenAI Codex was used for investigation, implementation, regression tests, and automated validation. --- graphene/utils/dataloader.py | 5 +- graphene/utils/tests/test_dataloader.py | 62 +++++++++++++++++++++++++ 2 files changed, 66 insertions(+), 1 deletion(-) diff --git a/graphene/utils/dataloader.py b/graphene/utils/dataloader.py index c75b9ee9e..015e3349d 100644 --- a/graphene/utils/dataloader.py +++ b/graphene/utils/dataloader.py @@ -261,6 +261,8 @@ async def dispatch_queue_batch(loader, queue): # Step through the values, resolving or rejecting each Future in the # loaded queue. for loaded, value in zip(queue, values): + if loaded.future.done(): + continue if isinstance(value, Exception): loaded.future.set_exception(value) else: @@ -277,4 +279,5 @@ def failed_dispatch(loader, queue, error): """ for loaded in queue: loader.clear(loaded.key) - loaded.future.set_exception(error) + if not loaded.future.done(): + loaded.future.set_exception(error) diff --git a/graphene/utils/tests/test_dataloader.py b/graphene/utils/tests/test_dataloader.py index 257f6b4db..8b0f341d0 100644 --- a/graphene/utils/tests/test_dataloader.py +++ b/graphene/utils/tests/test_dataloader.py @@ -450,3 +450,65 @@ async def do_resolve(x): a_loader, a_load_calls = id_loader(resolve=do_resolve) assert a_loader.clear("A1") == a_loader + + +@mark.asyncio +@mark.parametrize("failed", [False, True]) +@mark.parametrize("cancelled_index", [0, 1, 2]) +async def test_cancelled_load_does_not_interrupt_batch( + failed, cancelled_index, monkeypatch +): + from asyncio import gather + from graphene.utils.dataloader import dispatch_queue_batch, Loader + + error = RuntimeError("batch failure") + + async def resolve(keys): + if failed: + raise error + return keys + + loader = DataLoader(resolve) + # Dispatch explicitly so the regression deterministically exposes an + # invalid-state exception instead of waiting for an unresolved future. + monkeypatch.setattr( + loader, + "do_resolve_reject", + lambda key, future: loader._queue.append(Loader(key, future)), + ) + futures = [loader.load(key) for key in range(3)] + futures[cancelled_index].cancel() + await dispatch_queue_batch(loader, loader._queue) + for index, future in enumerate(futures): + if index == cancelled_index: + assert future.cancelled() + elif failed: + assert future.exception() is error + else: + assert future.result() == index + if failed: + assert not loader._cache + await gather(*futures, return_exceptions=True) + + +@mark.asyncio +async def test_cancelled_load_allows_scheduled_batch_to_finish(): + from asyncio import Event, wait_for + + started = Event() + release = Event() + + async def resolve(keys): + started.set() + await release.wait() + return keys + + loader = DataLoader(resolve) + cancelled = loader.load("cancelled") + surviving = loader.load("surviving") + await wait_for(started.wait(), 1) + cancelled.cancel() + release.set() + assert await wait_for(surviving, 1) == "surviving" + assert cancelled.cancelled() + assert await loader.load("surviving") == "surviving"