fix: compute mandatory populates even if not efficiently

Author: darkbasic

.changeset/twelve-lions-vanish.md

@@ -0,0 +1,5 @@
+---
+"mikro-orm-find-dataloader": minor
+---
+
+fix: compute mandatory populates even if not efficiently

packages/find/src/find.test.ts

@@ -124,9 +124,23 @@ describe("find", () => {
   it("should fetch books with the find dataloader", async () => {
     const authors = await em.fork().find(Author, {});
     const mock = mockLogger(orm);
-    const authorBooks = await Promise.all(
-      authors.map(async ({ id }) => await em.getRepository(Book).find({ author: id })),
-    );
+    const authorBooks = await Promise.all([
+      ...authors.map(async ({ id }) => await em.getRepository(Book).find({ author: id })),
+      // em.getRepository(Book).find({ author: { books: { author: 1 } } }),
+      // em.getRepository(Book).find({ title: "a", author: [1, { books: { author: 1 } }] }),
+      // em.getRepository(Book).find({ title: "a", author: { books: { author: 1 } } }),
+      // em.getRepository(Book).find({ title: "a", author: { books: { author: { name: "a" } } } }),
+      // em.getRepository(Book).find({ title: "a", author: { books: { title: "a" } } }),
+      // em.getRepository(Book).find({ title: "a", author: { books: 1 } }),
+      // em.getRepository(Book).find({ author: 1 }),
+      // em.getRepository(Book).find({ author: 1 }),
+      // em.getRepository(Book).find({ id: 2, title: "b", author: { id: 1, name: "a" } }),
+      // em.getRepository(Book).find({ author: { id: 1, name: "a" } }),
+      // em.getRepository(Book).find({ author: { id: 2 } }),
+      // em.getRepository(Book).find({ author: { id: 3 } }),
+      // em.getRepository(Book).find({ author: { name: "a" } }),
+    ]);
+    // console.log(mock.mock.calls);
     expect(authorBooks).toBeDefined();
     expect(authorBooks).toMatchSnapshot();
     expect(mock.mock.calls).toEqual([

packages/find/src/findDataloader.ts

@@ -124,7 +124,90 @@ export function groupInversedOrMappedKeysByEntity<T extends AnyEntity<T>>(
   return entitiesMap;
 }
 
+function allKeysArePK<K extends object>(
+  keys: Array<EntityKey<K>> | undefined,
+  primaryKeys: Array<EntityKey<K>>,
+): boolean {
+  if (keys == null) {
+    return false;
+  }
+  if (keys.length !== primaryKeys.length) {
+    return false;
+  }
+  for (const key of keys) {
+    if (!primaryKeys.includes(key)) {
+      return false;
+    }
+  }
+  return true;
+}
+
+// {id: 5, name: "a"} returns false because contains additional fields
+// Returns true for all PK formats including {id: 1} or {owner: 1, recipient: 2}
+function isPK<K extends object>(filter: FilterQueryDataloader<K>, meta: EntityMetadata<K>): boolean {
+  if (meta == null) {
+    return false;
+  }
+  if (meta.compositePK) {
+    // COMPOSITE
+    if (Array.isArray(filter)) {
+      // PK or PK[] or object[]
+      // [1, 2]
+      // [[1, 2], [3, 4]]
+      // [{owner: 1, recipient: 2}, {owner: 3, recipient: 4}]
+      // [{owner: 1, recipient: 2, sex: 0}, {owner: 3, recipient: 4, sex: 1}]
+      if (Utils.isPrimaryKey(filter, meta.compositePK)) {
+        // PK
+        return true;
+      }
+      if (Utils.isPrimaryKey(filter[0], meta.compositePK)) {
+        // PK[]
+        return true;
+      }
+      const keys = typeof filter[0] === "object" ? (Object.keys(filter[0]) as Array<EntityKey<K>>) : undefined;
+      if (allKeysArePK(keys, meta.primaryKeys)) {
+        // object is PK or PK[]
+        return true;
+      }
+    } else {
+      // object
+      // {owner: 1, recipient: 2, sex: 0}
+      const keys = typeof filter === "object" ? (Object.keys(filter) as Array<EntityKey<K>>) : undefined;
+      if (allKeysArePK(keys, meta.primaryKeys)) {
+        // object is PK
+        return true;
+      }
+    }
+  } else {
+    // NOT COMPOSITE
+    if (Array.isArray(filter)) {
+      // PK[]
+      // [1, 2]
+      // [{id: 1}, {id: 2}] NOT POSSIBLE FOR NON COMPOSITE
+      if (Utils.isPrimaryKey(filter[0])) {
+        return true;
+      }
+    } else {
+      // PK or object
+      // 1
+      // {id: [1, 2], sex: 0} or {id: 1, sex: 0}
+      if (Utils.isPrimaryKey(filter)) {
+        // PK
+        return true;
+      }
+      const keys =
+        typeof filter === "object" ? (Object.keys(filter) as [EntityKey<K>, ...Array<EntityKey<K>>]) : undefined;
+      if (keys?.length === 1 && keys[0] === meta.primaryKeys[0]) {
+        // object is PK
+        return true;
+      }
+    }
+  }
+  return false;
+}
+
 // Call this fn only if keyProp.targetMeta != null otherwise you will get false positives
+// Returns only PKs in short-hand format like 1 or [1, 1] not {id: 1} or {owner: 1, recipient: 2}
 function getPKs<K extends object>(
   filter: FilterQueryDataloader<K>,
   meta: EntityMetadata<K>,
@@ -180,20 +263,20 @@ COMPOSITE PK:
 const asc = (a: string, b: string): number => a.localeCompare(b);
 const notNull = (el: string | undefined): boolean => el != null;
 
-function getNewFiltersAndMapKeys<K extends object>(
+function getNewFiltersMapKeysAndMandatoryPopulate<K extends object>(
   cur: FilterQueryDataloader<K>,
   meta: EntityMetadata<K>,
   entityName: string,
-): Array<[FilterQueryDataloader<K>, string]>;
-function getNewFiltersAndMapKeys<K extends object>(
+): Array<[FilterQueryDataloader<K>, string, Set<string>?]>;
+function getNewFiltersMapKeysAndMandatoryPopulate<K extends object>(
   cur: FilterQueryDataloader<K>,
   meta: EntityMetadata<K>,
-): [FilterQueryDataloader<K>, string];
-function getNewFiltersAndMapKeys<K extends object>(
+): [FilterQueryDataloader<K>, string, string?];
+function getNewFiltersMapKeysAndMandatoryPopulate<K extends object>(
   cur: FilterQueryDataloader<K>,
   meta: EntityMetadata<K>,
   entityName?: string,
-): [FilterQueryDataloader<K>, string] | Array<[FilterQueryDataloader<K>, string]> {
+): [FilterQueryDataloader<K>, string, string?] | Array<[FilterQueryDataloader<K>, string, Set<string>?]> {
   const PKs = getPKs(cur, meta);
   if (PKs != null) {
     const res: [FilterQueryDataloader<K>, string] = [
@@ -206,6 +289,8 @@ function getNewFiltersAndMapKeys<K extends object>(
   } else {
     const newFilter: any = {};
     const keys: string[] = [];
+    const populateSet = new Set<string>();
+    let computedPopulate: string | undefined;
     if (Array.isArray(cur)) {
       // COMPOSITE PKs like [{owner: 1, recipient: 2}, {recipient: 4, owner: 3}]
       for (const key of meta.primaryKeys) {
@@ -223,27 +308,42 @@ function getNewFiltersAndMapKeys<K extends object>(
         // Using $or at the top level means that we can treat it as two separate queries and filter results from either of them
         if (key === "$or" && entityName != null) {
           return (value as Array<FilterQueryDataloader<K>>)
-            .map((el) => getNewFiltersAndMapKeys(el, meta, entityName))
+            .map((el) => getNewFiltersMapKeysAndMandatoryPopulate(el, meta, entityName))
             .flat();
         }
         const keyProp = meta.properties[key as EntityKey<K>];
         if (keyProp == null) {
           throw new Error(`Cannot find properties for ${key}`);
         }
         if (keyProp.targetMeta == null) {
+          // Our current key might lead to a scalar (thus we don't need to populate anything)
+          // or to an explicited PK like {id: 1, name: 'a'}
           newFilter[key] = Array.isArray(value) ? value : [value];
           keys.push(key);
         } else {
-          const [subFilter, subKey] = getNewFiltersAndMapKeys(value, keyProp.targetMeta);
+          // Our current key points to either a Reference or a Collection
+          const [subFilter, subKey, furtherPop] = getNewFiltersMapKeysAndMandatoryPopulate(value, keyProp.targetMeta);
           newFilter[key] = subFilter;
           keys.push(`${key}:${subKey}`);
+          // We need to populate all Collections and all the References
+          // where we further match non-PKs properties
+          if (keyProp.ref !== true || !isPK(value, keyProp.targetMeta)) {
+            computedPopulate = furtherPop == null ? `${key}` : `${key}.${furtherPop}`;
+            if (entityName != null) {
+              populateSet.add(computedPopulate);
+            } else {
+              // We return computedPopulate as the third element of the array
+            }
+          }
         }
       }
       const res: [FilterQueryDataloader<K>, string] = [
         newFilter,
         [entityName, `{${keys.sort(asc).join(",")}}`].filter(notNull).join("|"),
       ];
-      return entityName == null ? res : [res];
+      return entityName == null
+        ? [...res, computedPopulate]
+        : [[...res, populateSet.size === 0 ? undefined : populateSet]];
     }
   }
 }
@@ -259,7 +359,7 @@ function updateQueryFilter<K extends object, P extends string = never>(
   newQueryMap?: boolean,
 ): void {
   if (options?.populate != null && accOptions != null && accOptions.populate !== true) {
-    if (Array.isArray(options.populate) && options.populate[0] === "*") {
+    if (Array.isArray(options.populate) && options.populate.includes("*")) {
       accOptions.populate = true;
     } else if (Array.isArray(options.populate)) {
       if (accOptions.populate == null) {
@@ -276,14 +376,54 @@ function updateQueryFilter<K extends object, P extends string = never>(
       // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
       const curValue = (cur as Record<string, any[]>)[key]!;
       if (Array.isArray(value)) {
-        value.push(...curValue.reduce<any[]>((acc, cur) => acc.concat(cur), []));
+        // value.push(...curValue.reduce<any[]>((acc, cur) => acc.concat(cur), []));
+        value.push(...structuredClone(curValue));
       } else {
         updateQueryFilter([value], curValue);
       }
     }
   }
 }
 
+// The least amount of populate necessary to map the dataloader results to their original queries
+export function getMandatoryPopulate<K extends object>(
+  cur: FilterQueryDataloader<K>,
+  meta: EntityMetadata<K>,
+): string | undefined;
+export function getMandatoryPopulate<K extends object>(
+  cur: FilterQueryDataloader<K>,
+  meta: EntityMetadata<K>,
+  populate: Set<any>,
+): void;
+export function getMandatoryPopulate<K extends object>(
+  cur: FilterQueryDataloader<K>,
+  meta: EntityMetadata<K>,
+  populate?: Set<any>,
+): any {
+  for (const [key, value] of Object.entries(cur)) {
+    const keyProp = meta.properties[key as EntityKey<K>];
+    if (keyProp == null) {
+      throw new Error(`Cannot find properties for ${key}`);
+    }
+    // If our current key leads to scalar we don't need to populate anything
+    if (keyProp.targetMeta != null) {
+      // Our current key points to either a Reference or a Collection
+      // We need to populate all Collections
+      // We also need to populate References whenever we have to further match non-PKs properties
+      const PKs = getPKs(value, keyProp.targetMeta);
+      if (keyProp.ref !== true || PKs == null) {
+        const furtherPopulate = getMandatoryPopulate(value, keyProp.targetMeta);
+        const computedPopulate = furtherPopulate == null ? `${key}` : `${key}.${furtherPopulate}`;
+        if (populate != null) {
+          populate.add(computedPopulate);
+        } else {
+          return computedPopulate;
+        }
+      }
+    }
+  }
+}
+
 export interface DataloaderFind<K extends object, Hint extends string = never, Fields extends string = never> {
   entityName: string;
   meta: EntityMetadata<K>;
@@ -299,13 +439,13 @@ export function groupFindQueriesByOpts(
   const queriesMap = new Map<string, [FilterQueryDataloader<any>, { populate?: true | Set<any> }?]>();
   for (const dataloaderFind of dataloaderFinds) {
     const { entityName, meta, filter, options } = dataloaderFind;
-    const filtersAndKeys = getNewFiltersAndMapKeys(filter, meta, entityName);
+    const filtersAndKeys = getNewFiltersMapKeysAndMandatoryPopulate(filter, meta, entityName);
     dataloaderFind.filtersAndKeys = [];
-    filtersAndKeys.forEach(([newFilter, key]) => {
+    filtersAndKeys.forEach(([newFilter, key, mandatoryPopulate]) => {
       dataloaderFind.filtersAndKeys?.push({ key, newFilter });
       let queryMap = queriesMap.get(key);
       if (queryMap == null) {
-        queryMap = [structuredClone(newFilter), {}];
+        queryMap = [structuredClone(newFilter), { ...(mandatoryPopulate != null && { populate: mandatoryPopulate }) }];
         updateQueryFilter(queryMap, newFilter, options, true);
         queriesMap.set(key, queryMap);
       } else {
@@ -382,6 +522,9 @@ export function optsMapToQueries<Entity extends object>(
         populate: options.populate === true ? ["*"] : Array.from(options.populate),
       }),
     } satisfies Pick<FindOptions<any, any>, "populate">;
+    console.log("entityName", entityName);
+    console.log("filter", filter);
+    console.log("findOptions", findOptions);
     const entities = await em.find(entityName, filter, findOptions);
     return [key, entities];
   });