@nestjs/graphql でdataloaderを使う

GraphQLのN+1問題を解決するための策としてDataloaderがあります。今回はNestJSでDataloaderをProviderとして登録し、使用する方法を紹介します。

Note: Dataloaderの知識がある程度ある人向けに書いています。

本記事のゴール

このResolverクラスを

@Resolver("Query")
export class UserResolver {

  constructor(private userService: UserService) {}

  @ResolveProperty()
  public async user(@Args("id") id: number): Promise<User> {
    return this.userService.findById(id);
  }
}

次のように、DataloaderクラスをServiceの代わりにInjectできるようにします。

@Resolver("Query")
export class UserResolver {

  constructor(private userDataloader: UserDataloader) {}

  @ResolveProperty()
  public async user(@Args("id") id: number): Promise<User> {
    return this.userDataloader.load(id);
  }
}

本記事で覚えること

覚えることは1つです。「ProviderにScope.REQUESTを使ってリクエストごとにインスタンスを再生成する」ということを頭に入れておけばいいです。

Dataloaderクラスの作成

汎用性を高めるためにBaseDataloaderクラスを作成します。

import * as DataLoader from "dataloader";

export abstract class BaseDataloader<K, V> extends Object {
  protected dataloader: DataLoader<K, V> = new DataLoader<K, V>(this.batchLoad.bind(this));

  public clear(key: K): DataLoader<K, V> {
    return this.dataloader.clear(key);
  }

  public clearAll(): DataLoader<K, V> {
    return this.dataloader.clearAll();
  }

  public async load(key: K): Promise<V> {
    return this.dataloader.load(key);
  }

  public async loadMany(keys: K[]): Promise<V[]> {
    return this.dataloader.loadMany(keys);
  }

  public prime(key: K, value: V): DataLoader<K, V> {
    return this.dataloader.prime(key, value);
  }

  protected abstract async batchLoad(keys: K[]): Promise<(V | Error)[]>;
}

Dataloaderクラスを作成します。@Injectable({ scope: Scope.REQUEST }) というようにscopeを必ずScope.REQUESTにします。これでリクエストごとにDataloaderが初期化されます。

@Injectable({ scope: Scope.REQUEST })
export class UserDataloader extends BaseDataloader<number, User> {

  constructor(private readonly userService: UserService) {
    super();
  }

  protected async batchLoad(keys: number[]): Promise<(User | Error)[]> {
    return this.userService.findByIds(keys);
  }
}

あとはModuleにproviderとして登録するだけです。

@Module({
  imports: [TypeOrmModule.forFeature([UserRepository])],
  providers: [UserService, UserDataloader, UserResolver]
})
export class UserModule {}

これで、次のようにDataloaderを使用することが出来るようになります。

@Resolver("Query")
export class UserResolver {

  constructor(private userDataloader: UserDataloader) {}

  @ResolveProperty()
  public async user(@Args("id") id: number): Promise<User> {
    return this.userDataloader.load(id);
  }
}

Unitテスト時の注意

通常、TestingModuleからProviderを取得するときは module.get を使用すると思いますが、Scope.REQUESTを使用している場合は取得することができません。この場合、 userDataloaderはundefinedになります。代わりに module.resolve を使用すれば解決します。

const module = await Test.createTestingModule({
  providers: [UserResolver, UserDataloader, UserService],
}).compile();
// UserResolverのインスタンスは取得できるが userDataloaderがundefeind
// module.resolve<UserResolver>(UserResolver)
const resolver = await module.resolve<UserResolver>(UserResolver);

また、本記事で紹介したBaseDataloaderクラスを継承したものにspyOnを使用するときは

const dataloader = await module.resolve<UserDataloader>(UserDataloader);
jest.spyOn(dataloader.constructor.prototype, "load").mockResolvedValueOnce(mock);