fix(query-graphql): use DataLoader in ReferenceResolver to prevent N+1 queries (#389) (#391)

closes #389

Author: TriPSs

packages/query-graphql/__tests__/resolvers/reference.resolver.spec.ts

@@ -1,6 +1,7 @@
+import { ExecutionContext } from '@nestjs/common'
 import { Query, Resolver } from '@nestjs/graphql'
 import { ReferenceResolver, ReferenceResolverOpts } from '@ptc-org/nestjs-query-graphql'
-import { when } from 'ts-mockito'
+import { anything, when } from 'ts-mockito'
 
 import { createResolverFromNest, generateSchema, TestResolverDTO, TestService } from '../__fixtures__'
 
@@ -33,33 +34,48 @@ describe('ReferenceResolver', () => {
   })
 
   describe('#resolveReference', () => {
-    it('should call the service getById with the provided input', async () => {
+    const createContext = (): ExecutionContext => ({}) as unknown as ExecutionContext
+
+    it('should use DataLoader to resolve reference', async () => {
       const { resolver, mockService } = await createResolverFromNest(TestResolver)
+      const context = createContext()
       const id = 'id-1'
       const output: TestResolverDTO = {
         id,
         stringField: 'foo'
       }
-      when(mockService.getById(id)).thenResolve(output)
+
+      when(mockService.query(anything())).thenResolve([output])
+
       // @ts-ignore
       // eslint-disable-next-line @typescript-eslint/no-unsafe-call,@typescript-eslint/naming-convention
-      const result = await resolver.resolveReference({ __type: 'TestReference', id })
+      const result = await resolver.resolveReference({ __type: 'TestReference', id }, context)
       return expect(result).toEqual(output)
     })
 
     it('should reject if the id is not found', async () => {
-      const { resolver, mockService } = await createResolverFromNest(TestResolver)
-      const id = 'id-1'
-      const output: TestResolverDTO = {
-        id,
-        stringField: 'foo'
-      }
-      when(mockService.getById(id)).thenResolve(output)
+      const { resolver } = await createResolverFromNest(TestResolver)
+      const context = createContext()
+
       // @ts-ignore
       // eslint-disable-next-line @typescript-eslint/no-unsafe-call,@typescript-eslint/naming-convention
-      return expect(resolver.resolveReference({ __type: 'TestReference' })).rejects.toThrow(
+      return expect(resolver.resolveReference({ __type: 'TestReference' }, context)).rejects.toThrow(
         'Unable to resolve reference, missing required key id for TestResolverDTO'
       )
     })
+
+    it('should reject if entity is not found', async () => {
+      const { resolver, mockService } = await createResolverFromNest(TestResolver)
+      const context = createContext()
+      const id = 'id-not-found'
+
+      when(mockService.query(anything())).thenResolve([])
+
+      // @ts-ignore
+      // eslint-disable-next-line @typescript-eslint/no-unsafe-call,@typescript-eslint/naming-convention
+      return expect(resolver.resolveReference({ __type: 'TestReference', id }, context)).rejects.toThrow(
+        'Unable to find TestResolverDTO with id: id-not-found'
+      )
+    })
   })
 })

packages/query-graphql/src/loader/index.ts

@@ -3,3 +3,4 @@ export * from './count-relations.loader'
 export * from './dataloader.factory'
 export * from './find-relations.loader'
 export * from './query-relations.loader'
+export * from './reference.loader'

packages/query-graphql/src/loader/reference.loader.ts

@@ -0,0 +1,34 @@
+import { Class, Filter, QueryService } from '@ptc-org/nestjs-query-core'
+
+import { NestjsQueryDataloader } from './relations.loader'
+
+export type ReferenceArgs = {
+  id: string | number
+}
+
+export class ReferenceLoader<DTO> implements NestjsQueryDataloader<DTO, ReferenceArgs, DTO | undefined | Error> {
+  constructor(readonly DTOClass: Class<DTO>) {}
+
+  public createLoader(service: QueryService<DTO, unknown, unknown>) {
+    return async (args: ReadonlyArray<ReferenceArgs>): Promise<(DTO | undefined | Error)[]> => {
+      // Extract all unique IDs from the batch
+      const ids = args.map((arg) => arg.id)
+
+      // Use batch query to fetch all entities at once
+      const filter = { id: { in: ids } } as unknown as Filter<DTO>
+      const entities = await service.query({ filter })
+
+      // Create a map for fast lookup by ID
+      const entityMap = new Map<string | number, DTO>()
+      if (entities) {
+        entities.forEach((entity) => {
+          const id = (entity as Record<string, unknown>).id as string | number
+          entityMap.set(id, entity)
+        })
+      }
+
+      // Return results in the same order as requested
+      return args.map((arg) => entityMap.get(arg.id) || undefined)
+    }
+  }
+}

packages/query-graphql/src/resolvers/reference.resolver.ts

@@ -1,9 +1,12 @@
-import { BadRequestException } from '@nestjs/common'
-import { Resolver, ResolveReference } from '@nestjs/graphql'
+import { BadRequestException, ExecutionContext } from '@nestjs/common'
+import { Context, Resolver, ResolveReference } from '@nestjs/graphql'
 import { Class, QueryService } from '@ptc-org/nestjs-query-core'
 
 import { getDTONames } from '../common'
+import { InjectDataLoaderConfig } from '../decorators/inject-dataloader-config.decorator'
 import { RepresentationType } from '../federation'
+import { DataLoaderFactory, ReferenceLoader } from '../loader'
+import { DataLoaderOptions } from '../pipes/inject-data-loader-config.pipe'
 import { BaseServiceResolver, ResolverClass, ServiceResolver } from './resolver.interface'
 
 export interface ReferenceResolverOpts {
@@ -21,18 +24,37 @@ export const Referenceable =
       return BaseClass
     }
     const { key } = opts
+    const { baseName } = getDTONames(DTOClass)
+    const loaderName = `loadReference${baseName}`
+    const referenceLoader = new ReferenceLoader<DTO>(DTOClass)
 
     @Resolver(() => DTOClass, { isAbstract: true })
     class ResolveReferenceResolverBase extends BaseClass {
       @ResolveReference()
-      async resolveReference(representation: RepresentationType): Promise<DTO> {
+      async resolveReference(
+        representation: RepresentationType,
+        @Context() context: ExecutionContext,
+        @InjectDataLoaderConfig()
+        dataLoaderConfig?: DataLoaderOptions
+      ): Promise<DTO> {
         const id = representation[key]
         if (id === undefined) {
-          throw new BadRequestException(
-            `Unable to resolve reference, missing required key ${key} for ${getDTONames(DTOClass).baseName}`
-          )
+          throw new BadRequestException(`Unable to resolve reference, missing required key ${key} for ${baseName}`)
         }
-        return this.service.getById(representation[key] as string | number)
+
+        const loader = DataLoaderFactory.getOrCreateLoader(
+          context,
+          loaderName,
+          () => referenceLoader.createLoader(this.service),
+          dataLoaderConfig
+        )
+
+        const result = await loader.load({ id: id as string | number })
+        if (!result) {
+          throw new BadRequestException(`Unable to find ${baseName} with ${key}: ${String(id)}`)
+        }
+
+        return result
       }
     }