Skip to content

Commit 25afc86

Browse files
robrichardrobrichard
committed
add defer/stream support for subscriptions (#7)
# Conflicts: # src/subscription/subscribe.ts
1 parent 7058e4f commit 25afc86

File tree

4 files changed

+352
-20
lines changed

4 files changed

+352
-20
lines changed

‎src/execution/__tests__/flattenAsyncIterator-test.ts

Lines changed: 145 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,145 @@
1+
import { expect } from 'chai';
2+
import { describe, it } from 'mocha';
3+
4+
import { flattenAsyncIterator } from '../flattenAsyncIterator';
5+
6+
describe('flattenAsyncIterator', () => {
7+
it('does not modify an already flat async generator', async () => {
8+
async function* source() {
9+
yield await Promise.resolve(1);
10+
yield await Promise.resolve(2);
11+
yield await Promise.resolve(3);
12+
}
13+
14+
const result = flattenAsyncIterator(source());
15+
16+
expect(await result.next()).to.deep.equal({ value: 1, done: false });
17+
expect(await result.next()).to.deep.equal({ value: 2, done: false });
18+
expect(await result.next()).to.deep.equal({ value: 3, done: false });
19+
expect(await result.next()).to.deep.equal({
20+
value: undefined,
21+
done: true,
22+
});
23+
});
24+
25+
it('does not modify an already flat async iterator', async () => {
26+
const items = [1, 2, 3];
27+
28+
const iterator: any = {
29+
[Symbol.asyncIterator]() {
30+
return this;
31+
},
32+
next() {
33+
return Promise.resolve({
34+
done: items.length === 0,
35+
value: items.shift(),
36+
});
37+
},
38+
};
39+
40+
const result = flattenAsyncIterator(iterator);
41+
42+
expect(await result.next()).to.deep.equal({ value: 1, done: false });
43+
expect(await result.next()).to.deep.equal({ value: 2, done: false });
44+
expect(await result.next()).to.deep.equal({ value: 3, done: false });
45+
expect(await result.next()).to.deep.equal({
46+
value: undefined,
47+
done: true,
48+
});
49+
});
50+
51+
it('flatten nested async generators', async () => {
52+
async function* source() {
53+
yield await Promise.resolve(1);
54+
yield await Promise.resolve(2);
55+
yield await Promise.resolve(
56+
(async function* nested(): AsyncGenerator<number, void, void> {
57+
yield await Promise.resolve(2.1);
58+
yield await Promise.resolve(2.2);
59+
})(),
60+
);
61+
yield await Promise.resolve(3);
62+
}
63+
64+
const doubles = flattenAsyncIterator(source());
65+
66+
const result = [];
67+
for await (const x of doubles) {
68+
result.push(x);
69+
}
70+
expect(result).to.deep.equal([1, 2, 2.1, 2.2, 3]);
71+
});
72+
73+
it('allows returning early from a nested async generator', async () => {
74+
async function* source() {
75+
yield await Promise.resolve(1);
76+
yield await Promise.resolve(2);
77+
yield await Promise.resolve(
78+
(async function* nested(): AsyncGenerator<number, void, void> {
79+
yield await Promise.resolve(2.1); /* c8 ignore start */
80+
// Not reachable, early return
81+
yield await Promise.resolve(2.2);
82+
})(),
83+
);
84+
// Not reachable, early return
85+
yield await Promise.resolve(3);
86+
}
87+
/* c8 ignore stop */
88+
89+
const doubles = flattenAsyncIterator(source());
90+
91+
expect(await doubles.next()).to.deep.equal({ value: 1, done: false });
92+
expect(await doubles.next()).to.deep.equal({ value: 2, done: false });
93+
expect(await doubles.next()).to.deep.equal({ value: 2.1, done: false });
94+
95+
// Early return
96+
expect(await doubles.return()).to.deep.equal({
97+
value: undefined,
98+
done: true,
99+
});
100+
101+
// Subsequent next calls
102+
expect(await doubles.next()).to.deep.equal({
103+
value: undefined,
104+
done: true,
105+
});
106+
expect(await doubles.next()).to.deep.equal({
107+
value: undefined,
108+
done: true,
109+
});
110+
});
111+
112+
it('allows throwing errors from a nested async generator', async () => {
113+
async function* source() {
114+
yield await Promise.resolve(1);
115+
yield await Promise.resolve(2);
116+
yield await Promise.resolve(
117+
(async function* nested(): AsyncGenerator<number, void, void> {
118+
yield await Promise.resolve(2.1); /* c8 ignore start */
119+
// Not reachable, early return
120+
yield await Promise.resolve(2.2);
121+
})(),
122+
);
123+
// Not reachable, early return
124+
yield await Promise.resolve(3);
125+
}
126+
/* c8 ignore stop */
127+
128+
const doubles = flattenAsyncIterator(source());
129+
130+
expect(await doubles.next()).to.deep.equal({ value: 1, done: false });
131+
expect(await doubles.next()).to.deep.equal({ value: 2, done: false });
132+
expect(await doubles.next()).to.deep.equal({ value: 2.1, done: false });
133+
134+
// Throw error
135+
let caughtError;
136+
try {
137+
await doubles.throw('ouch'); /* c8 ignore start */
138+
// Not reachable, always throws
139+
/* c8 ignore stop */
140+
} catch (e) {
141+
caughtError = e;
142+
}
143+
expect(caughtError).to.equal('ouch');
144+
});
145+
});

‎src/execution/__tests__/subscribe-test.ts

Lines changed: 146 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -79,17 +79,22 @@ const emailSchema = new GraphQLSchema({
7979
}),
8080
});
8181

82-
function createSubscription(pubsub: SimplePubSub<Email>) {
82+
function createSubscription(
83+
pubsub: SimplePubSub<Email>,
84+
variableValues?: { readonly [variable: string]: unknown },
85+
) {
8386
const document = parse(`
84-
subscription ($priority: Int = 0) {
87+
subscription ($priority: Int = 0, $shouldDefer: Boolean = false) {
8588
importantEmail(priority: $priority) {
8689
email {
8790
from
8891
subject
8992
}
90-
inbox {
91-
unread
92-
total
93+
... @defer(if: $shouldDefer) {
94+
inbox {
95+
unread
96+
total
97+
}
9398
}
9499
}
95100
}
@@ -119,7 +124,12 @@ function createSubscription(pubsub: SimplePubSub<Email>) {
119124
}),
120125
};
121126

122-
return subscribe({ schema: emailSchema, document, rootValue: data });
127+
return subscribe({
128+
schema: emailSchema,
129+
document,
130+
rootValue: data,
131+
variableValues,
132+
});
123133
}
124134

125135
async function expectPromise(promise: Promise<unknown>) {
@@ -674,6 +684,136 @@ describe('Subscription Publish Phase', () => {
674684
});
675685
});
676686

687+
it('produces additional payloads for subscriptions with @defer', async () => {
688+
const pubsub = new SimplePubSub<Email>();
689+
const subscription = await createSubscription(pubsub, {
690+
shouldDefer: true,
691+
});
692+
assert(isAsyncIterable(subscription));
693+
// Wait for the next subscription payload.
694+
const payload = subscription.next();
695+
696+
// A new email arrives!
697+
expect(
698+
pubsub.emit({
699+
from: 'yuzhi@graphql.org',
700+
subject: 'Alright',
701+
message: 'Tests are good',
702+
unread: true,
703+
}),
704+
).to.equal(true);
705+
706+
// The previously waited on payload now has a value.
707+
expect(await payload).to.deep.equal({
708+
done: false,
709+
value: {
710+
data: {
711+
importantEmail: {
712+
email: {
713+
from: 'yuzhi@graphql.org',
714+
subject: 'Alright',
715+
},
716+
},
717+
},
718+
hasNext: true,
719+
},
720+
});
721+
722+
// Wait for the next payload from @defer
723+
expect(await subscription.next()).to.deep.equal({
724+
done: false,
725+
value: {
726+
data: {
727+
inbox: {
728+
unread: 1,
729+
total: 2,
730+
},
731+
},
732+
path: ['importantEmail'],
733+
hasNext: false,
734+
},
735+
});
736+
737+
// Another new email arrives, after all incrementally delivered payloads are received.
738+
expect(
739+
pubsub.emit({
740+
from: 'hyo@graphql.org',
741+
subject: 'Tools',
742+
message: 'I <3 making things',
743+
unread: true,
744+
}),
745+
).to.equal(true);
746+
747+
// The next waited on payload will have a value.
748+
expect(await subscription.next()).to.deep.equal({
749+
done: false,
750+
value: {
751+
data: {
752+
importantEmail: {
753+
email: {
754+
from: 'hyo@graphql.org',
755+
subject: 'Tools',
756+
},
757+
},
758+
},
759+
hasNext: true,
760+
},
761+
});
762+
763+
// Another new email arrives, before the incrementally delivered payloads from the last email was received.
764+
expect(
765+
pubsub.emit({
766+
from: 'adam@graphql.org',
767+
subject: 'Important',
768+
message: 'Read me please',
769+
unread: true,
770+
}),
771+
).to.equal(true);
772+
773+
// Deferred payload from previous event is received.
774+
expect(await subscription.next()).to.deep.equal({
775+
done: false,
776+
value: {
777+
data: {
778+
inbox: {
779+
unread: 2,
780+
total: 3,
781+
},
782+
},
783+
path: ['importantEmail'],
784+
hasNext: false,
785+
},
786+
});
787+
788+
// Next payload from last event
789+
expect(await subscription.next()).to.deep.equal({
790+
done: false,
791+
value: {
792+
data: {
793+
importantEmail: {
794+
email: {
795+
from: 'adam@graphql.org',
796+
subject: 'Important',
797+
},
798+
},
799+
},
800+
hasNext: true,
801+
},
802+
});
803+
804+
// The client disconnects before the deferred payload is consumed.
805+
expect(await subscription.return()).to.deep.equal({
806+
done: true,
807+
value: undefined,
808+
});
809+
810+
// Awaiting a subscription after closing it results in completed results.
811+
expect(await subscription.next()).to.deep.equal({
812+
done: true,
813+
value: undefined,
814+
});
815+
});
816+
677817
it('produces a payload when there are multiple events', async () => {
678818
const pubsub = new SimplePubSub<Email>();
679819
const subscription = await createSubscription(pubsub);

‎src/execution/flattenAsyncIterator.ts

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,50 @@
1+
import { isAsyncIterable } from '../jsutils/isAsyncIterable';
2+
3+
type AsyncIterableOrGenerator<T> =
4+
| AsyncGenerator<T, void, void>
5+
| AsyncIterable<T>;
6+
7+
/**
8+
* Given an AsyncIterable that could potentially yield other async iterators,
9+
* flatten all yielded results into a single AsyncIterable
10+
*/
11+
export function flattenAsyncIterator<T, AT>(
12+
iterable: AsyncIterableOrGenerator<T | AsyncIterableOrGenerator<AT>>,
13+
): AsyncGenerator<T | AT, void, void> {
14+
const iteratorMethod = iterable[Symbol.asyncIterator];
15+
const iterator: any = iteratorMethod.call(iterable);
16+
let iteratorStack: Array<AsyncIterator<T>> = [iterator];
17+
18+
async function next(): Promise<IteratorResult<T | AT, void>> {
19+
const currentIterator = iteratorStack[0];
20+
if (!currentIterator) {
21+
return { value: undefined, done: true };
22+
}
23+
const result = await currentIterator.next();
24+
if (result.done) {
25+
iteratorStack.shift();
26+
return next();
27+
} else if (isAsyncIterable(result.value)) {
28+
const childIterator = result.value[
29+
Symbol.asyncIterator
30+
]() as AsyncIterator<T>;
31+
iteratorStack.unshift(childIterator);
32+
return next();
33+
}
34+
return result;
35+
}
36+
return {
37+
next,
38+
return() {
39+
iteratorStack = [];
40+
return iterator.return();
41+
},
42+
throw(error?: unknown): Promise<IteratorResult<T | AT>> {
43+
iteratorStack = [];
44+
return iterator.throw(error);
45+
},
46+
[Symbol.asyncIterator]() {
47+
return this;
48+
},
49+
};
50+
}

0 commit comments

Comments
 (0)