diff --git a/src/execution/__tests__/stream-test.ts b/src/execution/__tests__/stream-test.ts index 8d04e84d39..a4bc92c451 100644 --- a/src/execution/__tests__/stream-test.ts +++ b/src/execution/__tests__/stream-test.ts @@ -133,6 +133,22 @@ async function complete(document: DocumentNode) { return result; } +async function completeAsync(document: DocumentNode, numCalls: number) { + const schema = new GraphQLSchema({ query }); + + const result = await execute({ schema, document, rootValue: {} }); + + invariant(isAsyncIterable(result)); + + const iterator = result[Symbol.asyncIterator](); + + const promises = []; + for (let i = 0; i < numCalls; i++) { + promises.push(iterator.next()); + } + return Promise.all(promises); +} + describe('Execute: stream directive', () => { it('Can stream a list field', async () => { const document = parse('{ scalarList @stream(initialCount: 1) }'); @@ -437,6 +453,58 @@ describe('Execute: stream directive', () => { }, ]); }); + it('Can stream a field that returns an async iterable', async () => { + const document = parse(` + query { + asyncIterableList @stream(initialCount: 2) { + name + id + } + } + `); + const result = await completeAsync(document, 4); + expectJSON(result).toDeepEqual([ + { + done: false, + value: { + data: { + asyncIterableList: [ + { + name: 'Luke', + id: '1', + }, + { + name: 'Han', + id: '2', + }, + ], + }, + hasNext: true, + }, + }, + { + done: false, + value: { + data: { + name: 'Leia', + id: '3', + }, + path: ['asyncIterableList', 2], + hasNext: true, + }, + }, + { + done: false, + value: { + hasNext: false, + }, + }, + { + done: true, + value: undefined, + }, + ]); + }); it('Handles error thrown in async iterable before initialCount is reached', async () => { const document = parse(` query { diff --git a/src/execution/execute.ts b/src/execution/execute.ts index 676534ad59..462d708905 100644 --- a/src/execution/execute.ts +++ b/src/execution/execute.ts @@ -1569,47 +1569,62 @@ export class Dispatcher { done: false, }); } - return new Promise<{ - promise: Promise>; - }>((resolve) => { + return new Promise((resolve) => { + let resolved = false; this._subsequentPayloads.forEach((promise) => { // eslint-disable-next-line @typescript-eslint/no-floating-promises - promise.then(() => { - // resolve with actual promise, not resolved value of promise so we can remove it from this._subsequentPayloads - resolve({ promise }); - }); - }); - }) - .then(({ promise }) => { - this._subsequentPayloads.splice( - this._subsequentPayloads.indexOf(promise), - 1, - ); - return promise; - }) - .then(({ value, done }) => { - if (done && this._subsequentPayloads.length === 0) { - // async iterable resolver just finished and no more pending payloads - return { - value: { - hasNext: false, - }, - done: false, + promise.then((payload) => { + if (resolved) { + return; + } + + resolved = true; + + if (this._subsequentPayloads.length === 0) { + // a different call to next has exhausted all payloads + resolve({ value: undefined, done: true }); + return; + } + + const index = this._subsequentPayloads.indexOf(promise); + + if (index === -1) { + // a different call to next has consumed this payload + resolve(this._race()); + return; + } + + this._subsequentPayloads.splice(index, 1); + + const { value, done } = payload; + + if (done && this._subsequentPayloads.length === 0) { + // async iterable resolver just finished and no more pending payloads + resolve({ + value: { + hasNext: false, + }, + done: false, + }); + return; + } else if (done) { + // async iterable resolver just finished but there are pending payloads + // return the next one + resolve(this._race()); + return; + } + + const returnValue: ExecutionPatchResult = { + ...value, + hasNext: this._subsequentPayloads.length > 0, }; - } else if (done) { - // async iterable resolver just finished but there are pending payloads - // return the next one - return this._race(); - } - const returnValue: ExecutionPatchResult = { - ...value, - hasNext: this._subsequentPayloads.length > 0, - }; - return { - value: returnValue, - done: false, - }; + resolve({ + value: returnValue, + done: false, + }); + }); }); + }); } _next(): Promise> {