feat(dataloader): address commments

packages/core/src/EntityManager.ts

@@ -52,22 +52,6 @@ import type { EntityComparator } from './utils/EntityComparator';
 import { OptimisticLockError, ValidationError } from './errors';
 import DataLoader from 'dataloader';
 
-function groupPrimaryKeysByEntity(
-  refs: readonly Reference<any>[],
-): Map<string, Set<Primary<any>>> {
-  const map = new Map<string, Set<Primary<any>>>();
-  for (const ref of refs) {
-    const className = helper(ref).__meta.className;
-    let primaryKeys = map.get(className);
-    if (primaryKeys == null) {
-      primaryKeys = new Set();
-      map.set(className, primaryKeys);
-    }
-    primaryKeys.add(helper(ref).getPrimaryKey() as Primary<any>);
-  }
-  return map;
-}
-
 /**
  * The EntityManager is the central access point to ORM functionality. It is a facade to all different ORM subsystems
  * such as UnitOfWork, Query Language and Repository API.
@@ -79,23 +63,14 @@ export class EntityManager<D extends IDatabaseDriver = IDatabaseDriver> {
   readonly _id = EntityManager.counter++;
   readonly global = false;
   readonly name = this.config.get('contextName');
+  readonly refLoader = new DataLoader(Utils.getRefBatchLoadFn(this));
   private readonly validator = new EntityValidator(this.config.get('strict'));
   private readonly repositoryMap: Dictionary<EntityRepository<any>> = {};
   private readonly entityLoader: EntityLoader = new EntityLoader(this);
   private readonly comparator = this.config.getComparator(this.metadata);
   private readonly entityFactory: EntityFactory = new EntityFactory(this);
   private readonly unitOfWork: UnitOfWork = new UnitOfWork(this);
   private readonly resultCache = this.config.getResultCacheAdapter();
-  private readonly refLoader = new DataLoader<Ref<any>, any>(async refs => {
-    const groupedIds = groupPrimaryKeysByEntity(refs);
-    const promises = Array.from(groupedIds).map(
-      async ([entity, ids]) =>
-        await this.getRepository(entity).find(Array.from(ids)),
-    );
-    await Promise.all(promises);
-    return await Promise.all(refs.map(async ref => await ref.load({ dataloader: false })));
-  });
-
   private filters: Dictionary<FilterDef> = {};
   private filterParams: Dictionary<Dictionary> = {};
   private transactionContext?: Transaction;

packages/core/src/entity/Reference.ts

@@ -114,7 +114,7 @@ export class Reference<T extends object> {
   async load<TT extends T, K extends keyof T = never, P extends string = never>(options?: LoadReferenceOptions<T, P> | K): Promise<Loaded<TT, P> | T[K]> {
     const opts: Dictionary = typeof options === 'object' ? options : { prop: options };
 
-    if (opts.dataloader === true) {
+    if (opts.dataloader) {
       return helper(this.entity).__em.refLoader.load(this);
     }
 

packages/core/src/utils/Utils.ts

@@ -21,11 +21,14 @@ import type {
   EntityProperty,
   IMetadataStorage,
   Primary,
+  Ref,
 } from '../typings';
 import { GroupOperator, PlainObject, QueryOperator, ReferenceKind } from '../enums';
 import type { Collection } from '../entity/Collection';
 import type { Platform } from '../platforms';
 import { helper } from '../entity/wrap';
+import { type EntityManager } from '../EntityManager';
+import type DataLoader from 'dataloader';
 
 export const ObjectBindingPattern = Symbol('ObjectBindingPattern');
 
@@ -1167,4 +1170,37 @@ export class Utils {
     return typeof value === 'object' && !!value && '__raw' in value;
   }
 
+  static groupPrimaryKeysByEntity(
+    refs: readonly Ref<any>[],
+  ): Map<string, Set<Primary<any>>> {
+    const map = new Map<string, Set<Primary<any>>>();
+    for (const ref of refs) {
+      const className = helper(ref).__meta.className;
+      let primaryKeys = map.get(className);
+      if (primaryKeys == null) {
+        primaryKeys = new Set();
+        map.set(className, primaryKeys);
+      }
+      primaryKeys.add(helper(ref).getPrimaryKey() as Primary<any>);
+    }
+    return map;
+  }
+
+  static getRefBatchLoadFn(em: EntityManager): DataLoader.BatchLoadFn<Ref<any>, any> {
+    return async (refs: readonly Ref<any>[]): Promise<ArrayLike<any | Error>> => {
+      const groupedIds = Utils.groupPrimaryKeysByEntity(refs);
+      const promises = Array.from(groupedIds).map(
+        ([entityName, ids]) =>
+          em.find(entityName, Array.from(ids)),
+      );
+      await Promise.all(promises);
+      /* Instead of assigning each find result to the original reference we use a shortcut
+        which takes advantage of the already existing Mikro-ORM caching mechanism:
+        when it calls ref.load it will automatically retrieve the entry the entity
+        from the cache (it will hit the cache because of the previous find query).
+        This trick won't be possible for collections where we have to map the results. */
+      return await Promise.all(refs.map(async ref => await ref.load({ dataloader: false })));
+    };
+  }
+
 }