From 8b29d99fa241d65bf880ec3efa4507995268df7b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Niccol=C3=B2=20Belli?= <niccolo.belli@linuxsystems.it>
Date: Wed, 10 May 2023 17:14:35 +0200
Subject: [PATCH] feat: collections dataloader

---
 packages/core/src/EntityManager.ts     |  1 +
 packages/core/src/entity/Collection.ts |  5 +-
 packages/core/src/utils/Utils.ts       | 73 +++++++++++++++++++++++++-
 3 files changed, 77 insertions(+), 2 deletions(-)

diff --git a/packages/core/src/EntityManager.ts b/packages/core/src/EntityManager.ts
index 983f8fc17fb2..16c790f0f676 100644
--- a/packages/core/src/EntityManager.ts
+++ b/packages/core/src/EntityManager.ts
@@ -64,6 +64,7 @@ export class EntityManager<D extends IDatabaseDriver = IDatabaseDriver> {
   readonly global = false;
   readonly name = this.config.get('contextName');
   readonly refLoader = new DataLoader(Utils.getRefBatchLoadFn(this));
+  readonly colLoader = new DataLoader(Utils.getColBatchLoadFn(this));
   private readonly validator = new EntityValidator(this.config.get('strict'));
   private readonly repositoryMap: Dictionary<EntityRepository<any>> = {};
   private readonly entityLoader: EntityLoader = new EntityLoader(this);
diff --git a/packages/core/src/entity/Collection.ts b/packages/core/src/entity/Collection.ts
index afc85fbd029e..41b10db9b6c7 100644
--- a/packages/core/src/entity/Collection.ts
+++ b/packages/core/src/entity/Collection.ts
@@ -67,7 +67,10 @@ export class Collection<T extends object, O extends object = object> extends Arr
   /**
    * Initializes the collection and returns the items
    */
-  async loadItems<TT extends T, P extends string = never>(options?: InitOptions<TT, P>): Promise<Loaded<TT, P>[]> {
+  async loadItems<TT extends T, P extends string = never>(options?: InitOptions<TT, P> & { dataloader?: boolean }): Promise<Loaded<TT, P>[]> {
+    if (options?.dataloader && !this.isInitialized()) {
+      return this.getEntityManager().colLoader.load(this);
+    }
     await this.load(options);
     return super.getItems() as Loaded<TT, P>[];
   }
diff --git a/packages/core/src/utils/Utils.ts b/packages/core/src/utils/Utils.ts
index 8b69805a2631..78efa46ff273 100644
--- a/packages/core/src/utils/Utils.ts
+++ b/packages/core/src/utils/Utils.ts
@@ -24,7 +24,7 @@ import type {
   Ref,
 } from '../typings';
 import { GroupOperator, PlainObject, QueryOperator, ReferenceKind } from '../enums';
-import type { Collection } from '../entity/Collection';
+import { Collection } from '../entity/Collection';
 import type { Platform } from '../platforms';
 import { helper } from '../entity/wrap';
 import { type EntityManager } from '../EntityManager';
@@ -1203,4 +1203,75 @@ export class Utils {
     };
   }
 
+  static groupInversedOrMappedKeysByEntity(
+    collections: readonly Collection<any>[],
+  ): Map<string, Map<string, Set<Primary<any>>>> {
+    const entitiesMap = new Map<string, Map<string, Set<Primary<any>>>>();
+    for (const col of collections) {
+      const className = col.property.type;
+      let propMap = entitiesMap.get(className);
+      if (propMap == null) {
+        propMap = new Map();
+        entitiesMap.set(className, propMap);
+      }
+      // Many to Many vs One to Many
+      const inversedProp: string | undefined = col.property.inversedBy ?? col.property.mappedBy;
+      if (inversedProp == null) {
+        throw new Error(
+          'Cannot find inversedBy or mappedBy prop: did you forget to set the inverse side of a many-to-many relationship?',
+        );
+      }
+      let primaryKeys = propMap.get(inversedProp);
+      if (primaryKeys == null) {
+        primaryKeys = new Set();
+        propMap.set(inversedProp, primaryKeys);
+      }
+      primaryKeys.add(helper(col.owner).getPrimaryKey());
+    }
+    return entitiesMap;
+  }
+
+  static getColBatchLoadFn(em: EntityManager): DataLoader.BatchLoadFn<Collection<any>, any> {
+    return async (collections: readonly Collection<any>[]) => {
+      const entitiesMap = Utils.groupInversedOrMappedKeysByEntity(collections);
+      const promises: Promise<any[]>[] = Array.from(entitiesMap.entries()).map(
+        async ([entityName, propMap]) =>
+          await em.getRepository(entityName).find(
+            {
+              $or: Array.from(propMap.entries()).map(([prop, pks]) => ({ [prop]: Array.from(pks) })),
+            },
+            {
+              // We need to populate collections in order to later retrieve the PKs from its items
+              populate: Array.from(propMap.keys()).filter(
+                prop => em.getMetadata().get(entityName).properties[prop]?.ref !== true,
+              ) as any,
+            },
+          ),
+      );
+      const results = (await Promise.all(promises)).flat();
+      return collections.map(collection =>
+        results.filter(result => {
+          // Entity matches
+          if (helper(result).__meta.className === collection.property.type) {
+            // Either inversedBy or mappedBy exist because we already checked in groupInversedOrMappedKeysByEntity
+            const refOrCol = result[(collection.property.inversedBy ?? collection.property.mappedBy)] as
+              | Ref<any>
+              | Collection<any>;
+            if (refOrCol instanceof Collection) {
+              for (const item of refOrCol.getItems()) {
+                if (helper(item).getPrimaryKey() === helper(collection.owner).getPrimaryKey()) {
+                  return true;
+                }
+              }
+            } else {
+              return helper(refOrCol).getPrimaryKey() === helper(collection.owner).getPrimaryKey();
+            }
+
+          }
+          return false;
+        }),
+      );
+    };
+  }
+
 }