Using GraphQL DataLoaders with NestJS

This post assumes familiarity with NestJS and GraphQL.

What we will be building

In this post we will build a simple GraphQL API in NestJS that enables getting a list of posts.

We will use the following GraphQL query:

query GetPosts {
  posts {
    id
    title
    body
    createdBy {
      id
      name
    }
  }
}

Creating NestJS application

nest new example-app

This will generate a new NestJS app with the following structure:

After removing what we don't need we are left with just app.module.ts and main.ts.

Adding users module

nest g module users

After generating module we will add user.entity.ts and users.service.ts:

user.entity.ts

export class User {
  id: number;
  name: string;
}

users.service.ts

import { Injectable } from '@nestjs/common';

import { delay } from '../util';
import { User } from './user.entity';

@Injectable()
export class UsersService {
  private users: User[] = [
    { id: 1, name: 'John' },
    { id: 2, name: 'Jane' },
    { id: 3, name: 'Alex' },
    { id: 4, name: 'Anna' },
  ];

  async getUsers() {
    console.log('Getting users...');
    await delay(3000);
    return this.users;
  }
}

Before we return users in getUsers method we simulate database latency with a delay of 3000ms.

Don't forget to add UsersService to exports array in users.module.ts

Adding posts module

Here we do pretty much the same we did in users module:

post.entity.ts

export class Post {
  id: string;
  title: string;
  body: string;
  userId: number;
}

posts.service.ts

import { Injectable } from '@nestjs/common';

import { delay } from '../util';
import { Post } from './post.entity';

@Injectable()
export class PostsService {
  private posts: Post[] = [
    { id: 'post-1', title: 'Post 1', body: 'Lorem 1', userId: 1 },
    { id: 'post-2', title: 'Post 2', body: 'Lorem 2', userId: 1 },
    { id: 'post-3', title: 'Post 3', body: 'Lorem 3', userId: 2 },
  ];

  async getPosts() {
    console.log('Getting posts...');
    await delay(3000);
    return this.posts;
  }
}

That should be enough for now when it comes to core logic. Now let's add GraphQL related code.

Adding GraphQL

We will be using code first aproach.

Installing packages

npm i @nestjs/graphql graphql-tools graphql apollo-server-express

Adding GraphQLModule to our AppModule:

import { Module } from '@nestjs/common';
import { GraphQLModule } from '@nestjs/graphql';
import { join } from 'path';

import { PostsModule } from './posts/posts.module';
import { UsersModule } from './users/users.module';

@Module({
  imports: [
    UsersModule,
    PostsModule,
    GraphQLModule.forRoot({
      autoSchemaFile: join(process.cwd(), 'src/schema.gql'),
    }),
  ],
  controllers: [],
  providers: [],
})
export class AppModule {}

By declaring autoSchemaFile property NestJS will generate GraphQL schema from types we declare in code. However since we haven't declared any when we run npm run start:dev we will get an error.
We will fix that error by declaring GraphQL types in our code. In order to do that we need to add some decorators to our entity classes:

user.entity.ts

import { Field, Int, ObjectType } from '@nestjs/graphql';

@ObjectType()
export class User {
  @Field(() => Int)
  id: number;

  @Field()
  name: string;
}

However this doesn't solve our problem since we are still getting an error. So adding a resolver should fix it:

users.resolver.ts

import { Query, Resolver } from '@nestjs/graphql';

import { User } from './user.entity';
import { UsersService } from './users.service';

@Resolver(User)
export class UsersResolver {
  constructor(private readonly usersService: UsersService) {}

  @Query(() => [User])
  getUsers() {
    return this.usersService.getUsers();
  }
}

Don't forget to add UsersResolver to providers array in users.module.ts

After adding UsersResolver the error goes away and we get a new file:

schema.gql

# ------------------------------------------------------
# THIS FILE WAS AUTOMATICALLY GENERATED (DO NOT MODIFY)
# ------------------------------------------------------

type User {
  id: Int!
  name: String!
}

type Query {
  getUsers: [User!]!
}

So let's test it out. Open GraphQL playground (usually on http://localhost:3000/graphql) and execute the following query:

query GetUsers {
  users {
    id
    name
  }
}

So after about 3 seconds we should get the following result:

{
  "data": {
    "users": [
      {
        "id": 1,
        "name": "John"
      },
      {
        "id": 2,
        "name": "Jane"
      },
      {
        "id": 3,
        "name": "Alex"
      },
      {
        "id": 4,
        "name": "Anna"
      }
    ]
  }
}

In the same way we will add decorators and resolver for posts:

post.entity.ts

import { Field, ObjectType } from '@nestjs/graphql';

@ObjectType()
export class Post {
  @Field()
  id: string;

  @Field()
  title: string;

  @Field()
  body: string;

  userId: number;
}

posts.resolver.ts

import { Query, Resolver } from '@nestjs/graphql';

import { Post } from './post.entity';
import { PostsService } from './posts.service';

@Resolver(Post)
export class PostsResolver {
  constructor(private readonly postsService: PostsService) {}

  @Query(() => [Post], { name: 'posts' })
  getPosts() {
    return this.postsService.getPosts();
  }
}

Adding relationships

So this is what GraphQL is all about: querying connected data.

We will now add createdBy field to post.entity.ts:

post.entity.ts

@Field(() => User)
createdBy?: User;

After this we should be able to run GetPosts query from the beginning of this post. However we get an error:

"Cannot return null for non-nullable field Post.createdBy."

In order to fix this we need to resolve createdBy field in posts.resolver.ts. We do that by adding the following methods:

posts.resolver.ts

@ResolveField('createdBy', () => User)
getCreatedBy(@Parent() post: Post) {
  const { userId } = post;
  return this.usersService.getUser(userId);
}

users.service.ts

async getUser(id: number) {
  console.log(`Getting user with id ${id}...`);
  await delay(1000);
  return this.users.find((user) => user.id === id);
}

We also have to export UsersService from UsersModule and then import UsersModule into PostsModule.

So now we can finally go ahead and run GetPosts query and we should get the following result:

{
  "data": {
    "posts": [
      {
        "id": "post-1",
        "title": "Post 1",
        "body": "Lorem 1",
        "createdBy": {
          "id": 1,
          "name": "John"
        }
      },
      {
        "id": "post-2",
        "title": "Post 2",
        "body": "Lorem 2",
        "createdBy": {
          "id": 1,
          "name": "John"
        }
      },
      {
        "id": "post-3",
        "title": "Post 3",
        "body": "Lorem 3",
        "createdBy": {
          "id": 2,
          "name": "Jane"
        }
      }
    ]
  }
}

So that took some time because of all those delays.
However if we check the console we should see the following:

Getting posts...
Getting user with id 1...
Getting user with id 1...
Getting user with id 2...

In a real world scenario all these lines would mean a separate query to the database. That is known as N+1 problem.

What this means is that for every post that first "query" returns we would have to make a separate query for it's creator even if all posts were created by the same person (as we can see above we are getting user with id 1 twice).

This is where DataLoader can help.

What is DataLoader

According to the official documentation:

DataLoader is a generic utility to be used as part of your application's data fetching layer to provide a simplified and consistent API over various remote data sources such as databases or web services via batching and caching.

Creating users loader

First we need to install it:

npm i dataloader

users.loader.ts

import * as DataLoader from 'dataloader';

import { mapFromArray } from '../util';
import { User } from './user.entity';
import { UsersService } from './users.service';

function createUsersLoader(usersService: UsersService) {
  return new DataLoader<number, User>(async (ids) => {
    const users = await usersService.getUsersByIds(ids);

    const usersMap = mapFromArray(users, (user) => user.id);

    return ids.map((id) => usersMap[id]);
  });
}

Let's explain what is happening here:

1. DataLoader constructor accepts a batching function as an argument. A batching function takes an array of ids (or keys) and returns a promise that resolves to an array of values. Important thing to note here is that those values must be in the exact same order as ids argument.

2. usersMap is a simple object where keys are user ids and values are actual users:
   

{
  1: {id: 1, name: "John"},
  ...
}

So let's see how this can be used:

const usersLoader = createUsersLoader(usersService);

const user1 = await usersLoader.load(1)
const user2 = await usersLoader.load(2);

This will actualy make one "database request" using that batching function we defined earlier and get users 1 and 2 at the same time.

How does this help in GraphQL

The basic idea is to create new users loader on every HTTP request so it can be used in multiple resolvers. In GraphQL a single request shares the same context object between resolvers so we should be able to "attach" our users loader to context and then use it in our resolvers.

Attaching values to GraphQL context

If we were using just Apollo Server we would attach values to context in the following way:

// Constructor
const server = new ApolloServer({
  typeDefs,
  resolvers,
  context: ({ req }) => ({
    authScope: getScope(req.headers.authorization)
  })
}));

// Example resolver
(parent, args, context, info) => {
  if(context.authScope !== ADMIN) throw new AuthenticationError('not admin');
  // Proceed
}

However in our NestJS application we don't explicitly instantiate ApolloServer so the context function should be declared when declaring GraphQLModule. In our case that's in app.module.ts:

GraphQLModule.forRoot({
  autoSchemaFile: join(process.cwd(), 'src/schema.gql'),
  context: () => ({
    randomValue: Math.random(),
  }),
}),

The next thing we should do is access context inside a resolver and in @nestjs/graphql there is a decorator for that:

posts.resolver.ts

@Query(() => [Post], { name: 'posts' })
getPosts(@Context() context: any) {
  console.log(context.randomValue);
  return this.postsService.getPosts();
}

@ResolveField('createdBy', () => User)
getCreatedBy(@Parent() post: Post, @Context() context: any {
  console.log(context.randomValue);
  const { userId } = post;
  return this.usersService.getUser(userId);
}

Now when we run GetPosts query we should see the following in the console:

0.858156868751532
Getting posts...
0.858156868751532
Getting user with id 1...
0.858156868751532
Getting user with id 1...
0.858156868751532
Getting user with id 2...

It's the same value for all resolvers and to prove that it is unique to each HTTP request we can just run the query again and check if randomValue is changed.

We can make this a bit nicer by passing a string to Context decorator:

@Query(() => [Post], { name: 'posts' })
getPosts(@Context('randomValue') randomValue: number) {
  console.log(randomValue);
  return this.postsService.getPosts();
}

@ResolveField('createdBy', () => User)
getCreatedBy(@Parent() post: Post, @Context('randomValue') randomValue: number) {
  console.log(randomValue);
  const { userId } = post;
  return this.usersService.getUser(userId);
}

Now that we've seen how to attach values to GraphQL context we can proceed and try to attach data loaders to it.

Attaching DataLoaders to GraphQL context

app.module.ts

GraphQLModule.forRoot({
  autoSchemaFile: join(process.cwd(), 'src/schema.gql'),
  context: () => ({
    randomValue: Math.random(),
    usersLoader: createUsersLoader(usersService),
  }),
}),

If we just try to add usersLoader as shown above we will get an error because usersService isn't defined. To solve this we need to change the definition for GraphQLModule to use forRootAsync method:

app.module.ts

GraphQLModule.forRootAsync({
  useFactory: (usersService: UsersService) => ({
    autoSchemaFile: join(process.cwd(), 'src/schema.gql'),
    context: () => ({
      randomValue: Math.random(),
      usersLoader: createUsersLoader(usersService),
    }),
  }),
}),

Now this may compile, but still won't actually work. We need to add inject property bellow useFactory:

useFactory: ...,
inject: [UsersService],

This will now throw an error so we need to somehow provide UsersService to GraphQLModule and we do that by importing UsersModule into GraphQLModule.

imports: [UsersModule],
useFactory: ...

With that we have now successfully attached usersLoader to GraphQL context object. Let's now see how to use it.

Using usersLoader inside a resolver

We can now go ahead and replace randomValue in our resolvers with usersLoader:

posts.resolver.ts

import { Context, Parent, Query, ResolveField, Resolver } from '@nestjs/graphql';
import * as DataLoader from 'dataloader';

import { User } from '../users/user.entity';
import { Post } from './post.entity';
import { PostsService } from './posts.service';

@Resolver(Post)
export class PostsResolver {
  constructor(private readonly postsService: PostsService) {}

  @Query(() => [Post], { name: 'posts' })
  getPosts() {
    return this.postsService.getPosts();
  }

  @ResolveField('createdBy', () => User)
  getCreatedBy(
    @Parent() post: Post,
    @Context('usersLoader') usersLoader: DataLoader<number, User>,
  ) {
    const { userId } = post;
    return usersLoader.load(userId);
  }
}

Now when we run GetPosts query the console output should look like this:

Getting posts...
Getting users with ids (1,2)

In a real world scenario this would mean just 2 database queries no matter the number of posts or users and that is how we solved the N+1 problem.

Conclusion

All this setup is a bit complex but the good thing is that it only needs to be done once and after that we can just add more loaders and use them in resolvers.

Full code is available on GitHub:
https://github.com/filipegeric/nestjs-graphql-dataloaders

Thanks for reading! :)