Skip to content

Commit 46889c0

Browse files
Allow multiple rootParams to be used with Dataloader child resolution (#65)
Allow child resolution via Dataloader to support multiple root params via specification of a batching parameter name
1 parent 59cba3f commit 46889c0

2 files changed

Lines changed: 162 additions & 29 deletions

File tree

src/service.js

Lines changed: 71 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -129,7 +129,7 @@ module.exports = function(mixinOptions) {
129129
*/
130130
createActionResolver(actionName, def = {}) {
131131
const {
132-
dataLoader = false,
132+
dataLoader: useDataLoader = false,
133133
nullIfError = false,
134134
params: staticParams = {},
135135
rootParams = {},
@@ -138,44 +138,63 @@ module.exports = function(mixinOptions) {
138138

139139
return async (root, args, context) => {
140140
try {
141-
if (dataLoader) {
141+
if (useDataLoader) {
142142
const dataLoaderMapKey = this.getDataLoaderMapKey(
143143
actionName,
144144
staticParams,
145145
args
146146
);
147-
const dataLoaderRootKey = rootKeys[0]; // for dataloader, use the first root key only
147+
// if a dataLoader batching parameter is specified, then all root params can be data loaded;
148+
// otherwise use only the primary rootParam
149+
const primaryDataLoaderRootKey = rootKeys[0]; // for dataloader, use the first root key only
150+
const dataLoaderBatchParam = this.dataLoaderBatchParams.get(actionName);
151+
const dataLoaderUseAllRootKeys = dataLoaderBatchParam != null;
148152

149153
// check to see if the DataLoader has already been added to the GraphQL context; if not then add it for subsequent use
150154
let dataLoader;
151155
if (context.dataLoaders.has(dataLoaderMapKey)) {
152156
dataLoader = context.dataLoaders.get(dataLoaderMapKey);
153157
} else {
154-
const batchedParamKey = rootParams[dataLoaderRootKey];
158+
const batchedParamKey =
159+
dataLoaderBatchParam || rootParams[primaryDataLoaderRootKey];
160+
155161
dataLoader = this.buildDataLoader(
156162
context.ctx,
157163
actionName,
158164
batchedParamKey,
159165
staticParams,
160-
args
166+
args,
167+
{ hashCacheKey: dataLoaderUseAllRootKeys } // must hash the cache key if not loading scalar
161168
);
162169
context.dataLoaders.set(dataLoaderMapKey, dataLoader);
163170
}
164171

165-
const rootValue = root && _.get(root, dataLoaderRootKey);
166-
if (rootValue == null) {
172+
let dataLoaderKey;
173+
if (dataLoaderUseAllRootKeys) {
174+
if (root && rootKeys) {
175+
dataLoaderKey = {};
176+
177+
rootKeys.forEach(key => {
178+
_.set(dataLoaderKey, rootParams[key], _.get(root, key));
179+
});
180+
}
181+
} else {
182+
dataLoaderKey = root && _.get(root, primaryDataLoaderRootKey);
183+
}
184+
185+
if (dataLoaderKey == null) {
167186
return null;
168187
}
169188

170-
return Array.isArray(rootValue)
171-
? await dataLoader.loadMany(rootValue)
172-
: await dataLoader.load(rootValue);
189+
return Array.isArray(dataLoaderKey)
190+
? await dataLoader.loadMany(dataLoaderKey)
191+
: await dataLoader.load(dataLoaderKey);
173192
} else {
174193
const params = {};
175194
if (root && rootKeys) {
176-
rootKeys.forEach(key =>
177-
_.set(params, rootParams[key], _.get(root, key))
178-
);
195+
rootKeys.forEach(key => {
196+
_.set(params, rootParams[key], _.get(root, key));
197+
});
179198
}
180199

181200
return await context.ctx.call(
@@ -223,21 +242,31 @@ module.exports = function(mixinOptions) {
223242
* @param {string} batchedParamKey - Parameter key to use for loaded values
224243
* @param {Object} staticParams - Static parameters to use in dataloader
225244
* @param {Object} args - Arguments passed to GraphQL child resolver
245+
* @param {Object} [options={}] - Optional arguments
246+
* @param {Boolean} [options.hashCacheKey=false] - Use a hash for the cacheKeyFn
226247
* @returns {DataLoader} Dataloader instance
227248
*/
228-
buildDataLoader(ctx, actionName, batchedParamKey, staticParams, args) {
249+
buildDataLoader(
250+
ctx,
251+
actionName,
252+
batchedParamKey,
253+
staticParams,
254+
args,
255+
{ hashCacheKey = false } = {}
256+
) {
229257
const batchLoadFn = keys => {
230258
const rootParams = { [batchedParamKey]: keys };
231259
return ctx.call(actionName, _.defaultsDeep({}, args, rootParams, staticParams));
232260
};
233261

234-
if (this.dataLoaderOptions.has(actionName)) {
235-
// use any specified options assigned to this action
236-
const options = this.dataLoaderOptions.get(actionName);
237-
return new DataLoader(batchLoadFn, options);
238-
}
262+
const dataLoaderOptions = this.dataLoaderOptions.get(actionName) || {};
263+
const cacheKeyFn = hashCacheKey && (key => hash(key));
264+
const options = {
265+
...(cacheKeyFn && { cacheKeyFn }),
266+
...dataLoaderOptions,
267+
};
239268

240-
return new DataLoader(batchLoadFn);
269+
return new DataLoader(batchLoadFn, options);
241270
},
242271

243272
/**
@@ -577,23 +606,39 @@ module.exports = function(mixinOptions) {
577606
*
578607
* @param {Object[]} services
579608
* @modifies {this.dataLoaderOptions}
609+
* @modifies {this.dataLoaderBatchParams}
580610
*/
581611
buildLoaderOptionMap(services) {
582612
this.dataLoaderOptions.clear(); // clear map before rebuilding
613+
this.dataLoaderBatchParams.clear(); // clear map before rebuilding
583614

584615
services.forEach(service => {
585616
Object.values(service.actions).forEach(action => {
586617
const { graphql: graphqlDefinition, name: actionName } = action;
587-
if (graphqlDefinition && graphqlDefinition.dataLoaderOptions) {
618+
if (
619+
graphqlDefinition &&
620+
(graphqlDefinition.dataLoaderOptions ||
621+
graphqlDefinition.dataLoaderBatchParam)
622+
) {
588623
const serviceName = this.getServiceName(service);
589624
const fullActionName = this.getResolverActionName(
590625
serviceName,
591626
actionName
592627
);
593-
this.dataLoaderOptions.set(
594-
fullActionName,
595-
graphqlDefinition.dataLoaderOptions
596-
);
628+
629+
if (graphqlDefinition.dataLoaderOptions) {
630+
this.dataLoaderOptions.set(
631+
fullActionName,
632+
graphqlDefinition.dataLoaderOptions
633+
);
634+
}
635+
636+
if (graphqlDefinition.dataLoaderBatchParam) {
637+
this.dataLoaderBatchParams.set(
638+
fullActionName,
639+
graphqlDefinition.dataLoaderBatchParam
640+
);
641+
}
597642
}
598643
});
599644
});
@@ -607,6 +652,7 @@ module.exports = function(mixinOptions) {
607652
this.pubsub = null;
608653
this.shouldUpdateGraphqlSchema = true;
609654
this.dataLoaderOptions = new Map();
655+
this.dataLoaderBatchParams = new Map();
610656

611657
const route = _.defaultsDeep(mixinOptions.routeOptions, {
612658
aliases: {

test/unit/service.spec.js

Lines changed: 91 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -210,7 +210,7 @@ describe("Test Service", () => {
210210

211211
it("should not invalidate schema when autoUpdateSchema is false", async () => {
212212
const { broker, svc, stop } = await startService({
213-
autoUpdateSchema: false
213+
autoUpdateSchema: false,
214214
});
215215
svc.invalidateGraphQLSchema = jest.fn();
216216

@@ -223,7 +223,7 @@ describe("Test Service", () => {
223223

224224
it("should not invalidate schema when autoUpdateSchema is true", async () => {
225225
const { broker, svc, stop } = await startService({
226-
autoUpdateSchema: true
226+
autoUpdateSchema: true,
227227
});
228228
svc.invalidateGraphQLSchema = jest.fn();
229229

@@ -548,6 +548,7 @@ describe("Test Service", () => {
548548

549549
beforeEach(() => {
550550
svc.dataLoaderOptions.clear();
551+
svc.dataLoaderBatchParams.clear();
551552
});
552553

553554
it("should return null if no rootValue", async () => {
@@ -637,6 +638,79 @@ describe("Test Service", () => {
637638
expect(ctx.call).toHaveBeenNthCalledWith(2, "users.resolve", { a: 5, id: [5] });
638639
});
639640

641+
it("should call the action via the loader using all root params", async () => {
642+
svc.dataLoaderBatchParams.set("users.resolve", "testBatchParam");
643+
const resolver = svc.createActionResolver("users.resolve", {
644+
rootParams: {
645+
authorId: "authorIdParam",
646+
testId: "testIdParam",
647+
},
648+
649+
dataLoader: true,
650+
});
651+
652+
const ctx = new Context(broker);
653+
ctx.call = jest.fn().mockResolvedValue(["res1", "res2", "res3"]);
654+
655+
const fakeRoot1 = { authorId: 1, testId: "foo" };
656+
const fakeRoot2 = { authorId: 2, testId: "bar" };
657+
const fakeRoot3 = { authorId: 5, testId: "baz" };
658+
659+
const fakeContext = { ctx, dataLoaders: new Map() };
660+
const res = await Promise.all([
661+
resolver(fakeRoot1, {}, fakeContext),
662+
resolver(fakeRoot2, {}, fakeContext),
663+
resolver(fakeRoot3, {}, fakeContext),
664+
]);
665+
666+
expect(res).toEqual(["res1", "res2", "res3"]);
667+
668+
expect(ctx.call).toHaveBeenCalledTimes(1);
669+
expect(ctx.call).toHaveBeenNthCalledWith(1, "users.resolve", {
670+
testBatchParam: [
671+
{ authorIdParam: 1, testIdParam: "foo" },
672+
{ authorIdParam: 2, testIdParam: "bar" },
673+
{ authorIdParam: 5, testIdParam: "baz" },
674+
],
675+
});
676+
});
677+
678+
it("should call the action via the loader using all root params while leveraging cache", async () => {
679+
svc.dataLoaderBatchParams.set("users.resolve", "testBatchParam");
680+
const resolver = svc.createActionResolver("users.resolve", {
681+
rootParams: {
682+
authorId: "authorIdParam",
683+
testId: "testIdParam",
684+
},
685+
686+
dataLoader: true,
687+
});
688+
689+
const ctx = new Context(broker);
690+
ctx.call = jest.fn().mockResolvedValue(["res1", "res2"]);
691+
692+
const fakeRoot1 = { authorId: 1, testId: "foo" };
693+
const fakeRoot2 = { authorId: 2, testId: "bar" };
694+
const fakeRoot3 = { authorId: 1, testId: "foo" }; // same as fakeRoot1
695+
696+
const fakeContext = { ctx, dataLoaders: new Map() };
697+
const res = await Promise.all([
698+
resolver(fakeRoot1, {}, fakeContext),
699+
resolver(fakeRoot2, {}, fakeContext),
700+
resolver(fakeRoot3, {}, fakeContext),
701+
]);
702+
703+
expect(res).toEqual(["res1", "res2", "res1"]);
704+
705+
expect(ctx.call).toHaveBeenCalledTimes(1);
706+
expect(ctx.call).toHaveBeenNthCalledWith(1, "users.resolve", {
707+
testBatchParam: [
708+
{ authorIdParam: 1, testIdParam: "foo" },
709+
{ authorIdParam: 2, testIdParam: "bar" },
710+
],
711+
});
712+
});
713+
640714
it("should reuse the loader for multiple calls with the same context", async () => {
641715
const resolver = svc.createActionResolver("users.resolve", {
642716
rootParams: {
@@ -1171,7 +1245,10 @@ describe("Test Service", () => {
11711245
actions: [
11721246
{
11731247
name: "test-action-1",
1174-
graphql: { dataLoaderOptions: { option1: "option-value-1" } },
1248+
graphql: {
1249+
dataLoaderOptions: { option1: "option-value-1" },
1250+
dataLoaderBatchParam: "batch-param-1",
1251+
},
11751252
},
11761253
{ name: "test-action-2" },
11771254
],
@@ -1182,7 +1259,10 @@ describe("Test Service", () => {
11821259
actions: [
11831260
{
11841261
name: "test-action-3",
1185-
graphql: { dataLoaderOptions: { option2: "option-value-2" } },
1262+
graphql: {
1263+
dataLoaderOptions: { option2: "option-value-2" },
1264+
dataLoaderBatchParam: "batch-param-2",
1265+
},
11861266
},
11871267
{ name: "test-action-4" },
11881268
],
@@ -1250,6 +1330,13 @@ describe("Test Service", () => {
12501330
])
12511331
);
12521332

1333+
expect(svc.dataLoaderBatchParams).toEqual(
1334+
new Map([
1335+
["test-svc-1.test-action-1", "batch-param-1"],
1336+
["v1.test-svc-2.test-action-3", "batch-param-2"],
1337+
])
1338+
);
1339+
12531340
// Test `context` method
12541341
const contextFn = ApolloServer.mock.calls[0][0].context;
12551342

0 commit comments

Comments
 (0)