Skip to content

Commit

Permalink
Return underlying AsyncIterators when execute result is returned (#2843)
Browse files Browse the repository at this point in the history
  • Loading branch information
robrichard committed Feb 16, 2021
1 parent 4c262a2 commit 24330c0
Show file tree
Hide file tree
Showing 2 changed files with 167 additions and 0 deletions.
141 changes: 141 additions & 0 deletions src/execution/__tests__/stream-test.js
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
import { expect } from 'chai';
import { describe, it } from 'mocha';

import { invariant } from '../../jsutils/invariant';
import { isAsyncIterable } from '../../jsutils/isAsyncIterable';
import { parse } from '../../language/parser';

Expand Down Expand Up @@ -72,6 +73,36 @@ const query = new GraphQLObjectType({
yield {};
},
},
asyncIterableListDelayed: {
type: new GraphQLList(friendType),
async *resolve() {
for (const friend of friends) {
// pause an additional ms before yielding to allow time
// for tests to return or throw before next value is processed.
// eslint-disable-next-line no-await-in-loop
await new Promise((r) => setTimeout(r, 1));
yield friend;
}
},
},
asyncIterableListNoReturn: {
type: new GraphQLList(friendType),
resolve() {
let i = 0;
return {
[Symbol.asyncIterator]: () => ({
async next() {
const friend = friends[i++];
if (friend) {
await new Promise((r) => setTimeout(r, 1));
return { value: friend, done: false };
}
return { value: undefined, done: true };
},
}),
};
},
},
asyncIterableListDelayedClose: {
type: new GraphQLList(friendType),
async *resolve() {
Expand Down Expand Up @@ -649,4 +680,114 @@ describe('Execute: stream directive', () => {
},
]);
});
it('Returns underlying async iterables when dispatcher is returned', async () => {
const document = parse(`
query {
asyncIterableListDelayed @stream(initialCount: 1) {
name
id
}
}
`);
const schema = new GraphQLSchema({ query });

const executeResult = await execute({ schema, document, rootValue: {} });
invariant(isAsyncIterable(executeResult));

const result1 = await executeResult.next();
expect(result1).to.deep.equal({
done: false,
value: {
data: {
asyncIterableListDelayed: [
{
id: '1',
name: 'Luke',
},
],
},
hasNext: true,
},
});

executeResult.return();

// this result had started processing before return was called
const result2 = await executeResult.next();
expect(result2).to.deep.equal({
done: false,
value: {
data: {
id: '2',
name: 'Han',
},
hasNext: true,
path: ['asyncIterableListDelayed', 1],
},
});

// third result is not returned because async iterator has returned
const result3 = await executeResult.next();
expect(result3).to.deep.equal({
done: false,
value: {
hasNext: false,
},
});
});
it('Can return async iterable when underlying iterable does not have a return method', async () => {
const document = parse(`
query {
asyncIterableListNoReturn @stream(initialCount: 1) {
name
id
}
}
`);
const schema = new GraphQLSchema({ query });

const executeResult = await execute({ schema, document, rootValue: {} });
invariant(isAsyncIterable(executeResult));

const result1 = await executeResult.next();
expect(result1).to.deep.equal({
done: false,
value: {
data: {
asyncIterableListNoReturn: [
{
id: '1',
name: 'Luke',
},
],
},
hasNext: true,
},
});

executeResult.return();

// this result had started processing before return was called
const result2 = await executeResult.next();
expect(result2).to.deep.equal({
done: false,
value: {
data: {
id: '2',
name: 'Han',
},
hasNext: true,
path: ['asyncIterableListNoReturn', 1],
},
});

// third result is not returned because async iterator has returned
const result3 = await executeResult.next();
expect(result3).to.deep.equal({
done: false,
value: {
hasNext: false,
},
});
});
});
26 changes: 26 additions & 0 deletions src/execution/execute.js
Original file line number Diff line number Diff line change
Expand Up @@ -1608,11 +1608,15 @@ type DispatcherResult = {|
*/
export class Dispatcher {
_subsequentPayloads: Array<Promise<IteratorResult<DispatcherResult, void>>>;
_iterators: Array<AsyncIterator<mixed>>;
_isDone: boolean;
_initialResult: ?ExecutionResult;
_hasReturnedInitialResult: boolean;

constructor() {
this._subsequentPayloads = [];
this._iterators = [];
this._isDone = false;
this._hasReturnedInitialResult = false;
}

Expand Down Expand Up @@ -1681,13 +1685,16 @@ export class Dispatcher {
itemType: GraphQLOutputType,
): void {
const subsequentPayloads = this._subsequentPayloads;
const iterators = this._iterators;
iterators.push(iterator);
function next(index) {
const fieldPath = addPath(path, index);
const patchErrors = [];
subsequentPayloads.push(
iterator.next().then(
({ value: data, done }) => {
if (done) {
iterators.splice(iterators.indexOf(iterator), 1);
return { value: undefined, done: true };
}

Expand Down Expand Up @@ -1758,6 +1765,14 @@ export class Dispatcher {
}

_race(): Promise<IteratorResult<ExecutionPatchResult, void>> {
if (this._isDone) {
return Promise.resolve({
value: {
hasNext: false,
},
done: false,
});
}
return new Promise((resolve) => {
this._subsequentPayloads.forEach((promise) => {
promise.then(() => {
Expand Down Expand Up @@ -1814,13 +1829,24 @@ export class Dispatcher {
return this._race();
}

_return(): Promise<IteratorResult<AsyncExecutionResult, void>> {
return Promise.all(
// $FlowFixMe[prop-missing]
this._iterators.map((iterator) => iterator.return?.()),
).then(() => {
this._isDone = true;
return { value: undefined, done: true };
});
}

get(initialResult: ExecutionResult): AsyncIterable<AsyncExecutionResult> {
this._initialResult = initialResult;
return ({
[Symbol.asyncIterator]() {
return this;
},
next: () => this._next(),
return: () => this._return(),
}: any);
}
}
Expand Down

0 comments on commit 24330c0

Please sign in to comment.