Dataloader helpers for apollo server

1. README.md

This is a set of functions that could be used as helpers for the Dataloader (https://github.com/graphql/dataloader). The helpers are designed to handle two cases:

1. Batching fetches when each key corresponds to a single value (think userId -> User)
2. Batching fetches when each key corresponds to multiple values (think userId -> userPosts)

One item per key

For the first case, there is the mapKey function.

It wraps the batch loading function to make sure it always returns an array the same length as the keys array and with items in the same order as the keys.

• The first arg is the actual batch loading function
• The second arg is a function to extract the id from an item returned by the batch loading function

Many items per key

For the second case, there is the groupByKey function.

It wraps the batch loading function to group items by key.

• The first arg is the actual batch loading function. It returns all the items mixed in a single.
• The second arg is a function to extract the related id from each item. The items will be grouped based on this extracted id.

dataloader.js

const { ApolloServer, gql } = require('apollo-server');
const db = require('./models');
const DataLoader = require('dataloader');

// user loader
const user = loaderWithContext(
  mapKey(keys => db.User.findAll({ where: { id: keys } }), user => user.id),
);

// user posts loader
const postsByUser = loaderWithContext(
  groupByKey(
    userIds => db.Post.findAll({ where: { userId: userIds } }),
    post => post.userId,
  ),
);

const resolvers = {
  Query: Query(),
  Mutation: Mutation(),
  Post: {
    ...timestamps(),
    user: (post, _, context) => user(context).load(post.userId),
  },
  User: {
    ...timestamps(),
    posts: (user, _, context) => postsByUser(context).load(user.id),
  },
};

// Thanks to https://github.com/graphql/dataloader/issues/158 for this!
function loaderWithContext(fn) {
  const store = new WeakMap();
  return context => {
    let loader = store.get(context);
    if (!loader) {
      loader = new DataLoader(fn);
      store.set(context, loader);
    }
    return loader;
  };
}

/**
 * Dataloader helper. It will make sure Dataloader's conditions are met correctly (https://github.com/graphql/dataloader#batch-function)
 * 1. It will make the array of values the same length as the array of keys
 * 2. It will make sure to order the values the same as the keys, returning null for any values not found
 *
 * @param {*} fn your batch loading function
 * @param {*} getKey function that given an item, it returns the key it corresponds to
 * @returns a batch loading function that meets conditions 1 and 2
 */
function mapKey(fn, getKey = it => it.id) {
  return keys => {
    return fn(keys).then(items => {
      // construct a map of key -> item
      const map = new Map();
      items.forEach(it => map.set(getKey(it), it));

      // map each key to its item, or null
      return keys.map(id => (map.has(id) ? map.get(id) : null));
    });
  };
}

/**
 * Dataloader helper. It will group a list of values according to a given key, and return them in the corresponding order
 * according to the array of keys given
 * @param {*} fn
 * @param {*} getKey
 */
function groupByKey(fn, getKey) {
  return keys => {
    return fn(keys).then(mixedItems => {
      // group items by an id
      const itemsById = arrayGroupByKey(mixedItems, getKey);

      // return the items corresponding to the key, or an empty array
      return keys.map(key => itemsById.get(key) || []);
    });
  };
}

function contextLoader(resolver, fn) {
  const store = new WeakMap();
  return (arg, params, ctx) => {
    let loader = store.get(ctx);
    if (!loader) {
      loader = new DataLoader(fn);
      store.set(ctx, loader);
    }
    return resolver(arg, params, ctx, loader);
  };
}

const typeDefs = gql`
  type User {
    id: Int!
    firstName: String
    lastName: String
    email: String
    createdAt: String!
    updatedAt: String!
    posts: [Post!]
  }

  type Post {
    id: Int!
    content: String
    userId: Int!
    createdAt: String!
    updatedAt: String!
    user: User!
  }

  type Query {
    users: [User!]!
    posts: [Post!]!
  }

  type Mutation {
    addPost(userId: Int!): AddPostPayload!
    addUser(userInput: UserInput!): User!
  }
  type AddPostPayload {
    post: Post!
  }

  input UserInput {
    firstName: String!
    lastName: String!
    email: String!
  }
`;

const server = new ApolloServer({
  typeDefs,
  resolvers,
  // set an empty object of loaders for each new request
  context: async () => ({ loaders: new Map() }),
});

server.listen().then(({ url }) => {
  console.log(`🚀  Server ready at ${url}`);
});
process.on('exit', server.stop);

function timestamps() {
  return {
    createdAt: item => item.createdAt.toISOString(),
    updatedAt: item => item.updatedAt.toISOString(),
  };
}

/**
 * Group an array's items by a specific value, given by getKey
 * @param array the items to group
 * @param getKey a function that given an item, returns a value to group it by
 * @returns a map of key -> item[]
 */
function arrayGroupByKey(array, getKey) {
  const itemsById = new Map();
  array.forEach(item => {
    const id = getKey(item);
    let items = itemsById.get(id);
    if (!items) {
      items = [];
      itemsById.set(id, items);
    }
    items.push(item);
  });
  return itemsById;
}

function Query() {
  return {
    users: () => db.User.findAll().then(a => a.map(it => it.get())),
    posts: () => db.Post.findAll().then(a => a.map(it => it.get())),
  };
}

function Mutation() {
  return {
    addUser: async (_, { userInput }) =>
      db.User.create(userInput).then(a => a.get()),
    addPost: async (_, { userId }) => {
      const user = await db.User.findByPk(userId).then(a => a && a.get());
      if (!user) {
        throw new Error('No user for id ' + userId);
      }
      const post = await db.Post.create({
        content: `This is a post for the user ${user.id}: ${user.firstName} ${
          user.lastName
        }`,
        userId,
      }).then(a => a.get());
      return {
        post: post,
      };
    },
  };
}