fix(reference-loader): make key dynamic in createLoader and loader calls

Author: abrenoch

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

@@ -4,43 +4,43 @@ import { Class, Filter, QueryService } from '@ptc-org/nestjs-query-core'
 import { NestjsQueryDataloader } from './relations.loader'
 
 export type ReferenceArgs = {
-  id: string | number
+  [key: string]: string | number
 }
 
 export class ReferenceLoader<DTO> implements NestjsQueryDataloader<DTO, ReferenceArgs, DTO | undefined | Error> {
   private readonly logger = new Logger(ReferenceLoader.name)
 
   constructor(readonly DTOClass: Class<DTO>) {}
 
-  public createLoader(service: QueryService<DTO, unknown, unknown>) {
+  public createLoader(service: QueryService<DTO, unknown, unknown>, key = 'id') {
     return async (args: ReadonlyArray<ReferenceArgs>): Promise<(DTO | undefined | Error)[]> => {
       // Deduplicate IDs while preserving original order mapping
-      const uniqueIds = [...new Set(args.map((arg) => arg.id))]
+      const uniqueIds = [...new Set(args.map((arg) => arg[key]))]
 
       try {
-        const filter = { id: { in: uniqueIds } } as unknown as Filter<DTO>
+        const filter = { [key]: { in: uniqueIds } } as unknown as Filter<DTO>
         const entities = await service.query({ filter })
 
         // Create entity map for efficient lookup
         // Use string keys because representation.id from federation may differ in type from entity.id
         const entityMap = new Map<string, DTO>()
         entities?.forEach((entity) => {
-          const id = (entity as Record<string, unknown>).id
+          const id = (entity as Record<string, unknown>)[key]
           entityMap.set(String(id), entity)
         })
 
         // Return results in the same order as requested
         const results = args.map((arg) => {
-          const entity = entityMap.get(String(arg.id))
+          const entity = entityMap.get(String(arg[key]))
           if (!entity) {
-            this.logger.warn(`Entity not found for ID: ${arg.id} in ${this.DTOClass.name}`)
+            this.logger.warn(`Entity not found for ${key}: ${arg[key]} in ${this.DTOClass.name}`)
           }
           return entity || undefined
         })
 
         return results
       } catch (error) {
-        this.logger.error(`Batch query failed for ${this.DTOClass.name}:`, error)
+        this.logger.error(`Batch query failed for ${this.DTOClass.name} with key ${key}:`, error)
         // Return error for each request, let DataLoader handle it
         return args.map(() => (error instanceof Error ? error : new Error(String(error))))
       }

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

@@ -56,7 +56,7 @@ export const Referenceable =
         const loader = DataLoaderFactory.getOrCreateLoader(
           context,
           uniqueLoaderName,
-          () => referenceLoader.createLoader(this.service),
+          () => referenceLoader.createLoader(this.service, key),
           {
             // Ensure batching is enabled for performance
             batch: true,
@@ -66,7 +66,7 @@ export const Referenceable =
           }
         )
 
-        const result = await loader.load({ id: id as string | number })
+        const result = await loader.load({ [key]: id as string | number })
 
         if (!result) {
           this.logger.error(`Unable to find ${baseName} with ${key}: ${String(id)}`)