mikro-orm-find-dataloader/src/findDataloader.ts

Additional dataloaders for the MikroORM EntityManager find/findOne/etc methods.

import {
  Utils,
  Collection,
  helper,
  type FindOptions,
  type AnyEntity,
  type Reference,
  type Primary,
  type EntityMetadata,
  type EntityKey,
  type FilterItemValue,
  type Scalar,
  type ExpandProperty,
  type ExpandQuery,
  type ExpandScalar,
  type EntityProps,
  type EntityManager,
  type EntityName,
} from "@mikro-orm/core";
import { type PartialBy } from "./types";
import type DataLoader from "dataloader";

/* eslint-disable @typescript-eslint/ban-types */
/* eslint-disable @typescript-eslint/array-type */

export interface OperatorMapDataloader<T> {
  // $and?: ExpandQuery<T>[];
  $or?: Array<ExpandQuery<T>>;
  // $eq?: ExpandScalar<T> | ExpandScalar<T>[];
  // $ne?: ExpandScalar<T>;
  // $in?: ExpandScalar<T>[];
  // $nin?: ExpandScalar<T>[];
  // $not?: ExpandQuery<T>;
  // $gt?: ExpandScalar<T>;
  // $gte?: ExpandScalar<T>;
  // $lt?: ExpandScalar<T>;
  // $lte?: ExpandScalar<T>;
  // $like?: string;
  // $re?: string;
  // $ilike?: string;
  // $fulltext?: string;
  // $overlap?: string[];
  // $contains?: string[];
  // $contained?: string[];
  // $exists?: boolean;
}

export type FilterValueDataloader<T> =
  /* OperatorMapDataloader<FilterItemValue<T>> | */
  FilterItemValue<T> | FilterItemValue<T>[] | null;

export type QueryDataloader<T> = T extends object
  ? T extends Scalar
    ? never
    : FilterQueryDataloader<T>
  : FilterValueDataloader<T>;

export type FilterObjectDataloader<T> = {
  -readonly [K in EntityKey<T>]?:
    | QueryDataloader<ExpandProperty<T[K]>>
    | FilterValueDataloader<ExpandProperty<T[K]>>
    | null;
};

export type Compute<T> = {
  [K in keyof T]: T[K];
} & {};

export type ObjectQueryDataloader<T> = Compute<OperatorMapDataloader<T> & FilterObjectDataloader<T>>;

// FilterQuery<T>
export type FilterQueryDataloader<T extends object> =
  | ObjectQueryDataloader<T>
  | NonNullable<ExpandScalar<Primary<T>>> // Just 5 (or [5, 7] for composite keys). Currently not supported, we do {id: number} instead. Should be easy to add.
  // Accepts {id: 5} or any scalar like {name: "abc"}, IdentifiedReference (because it extends {id: 5}) but not just 5 nor {location: IdentifiedReference} (don't know why).
  // OperatorMap must be cut down to just a couple.
  | NonNullable<EntityProps<T> & OperatorMapDataloader<T>>
  | FilterQueryDataloader<T>[];

/* eslint-enable @typescript-eslint/ban-types */
/* eslint-enable @typescript-eslint/array-type */

export function groupPrimaryKeysByEntity<T extends AnyEntity<T>>(
  refs: Array<Reference<T>>,
): Map<string, Set<Primary<T>>> {
  const map = new Map<string, Set<Primary<T>>>();
  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<T>);
  }
  return map;
}

export function groupInversedOrMappedKeysByEntity<T extends AnyEntity<T>>(
  collections: Array<Collection<T, AnyEntity>>,
): Map<string, Map<string, Set<Primary<T>>>> {
  const entitiesMap = new Map<string, Map<string, Set<Primary<T>>>>();
  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() as Primary<T>);
  }
  return entitiesMap;
}

// Call this fn only if keyProp.targetMeta != null otherwise you will get false positives
function getPKs<K extends object>(
  filter: FilterQueryDataloader<K>,
  meta: EntityMetadata<K>,
): Array<Primary<K>> | undefined {
  if (meta.compositePK) {
    // COMPOSITE
    if (Array.isArray(filter)) {
      // PK or PK[] or object[]
      if (Utils.isPrimaryKey(filter, meta.compositePK)) {
        // PK
        return [filter as Primary<K>];
      }
      if (Utils.isPrimaryKey(filter[0], meta.compositePK)) {
        // PK[]
        return filter as Array<Primary<K>>;
      }
    }
  } else {
    // NOT COMPOSITE
    if (Array.isArray(filter)) {
      // PK[] or object[]
      if (Utils.isPrimaryKey(filter[0])) {
        return filter as Array<Primary<K>>;
      }
    } else {
      // PK or object
      if (Utils.isPrimaryKey(filter)) {
        // PK
        return [filter as Primary<K>];
      }
    }
  }
}

/*
NOT COMPOSITE:                                                                     NEW QUERY                                                      MAP KEY (props in alphabetical order)
  1                                                                             -> {id: [1]}                                                      {id}
  [1, 2]                                                                        -> {id: [1, 2]}                                                   {id}
  {id: 1, sex: 0}                                                               -> {id: [1], sex: [0]}                                            {id,sex}
  {id: [1, 2], sex: 0}                                                          -> {id: [1, 2], sex: [0]}                                         {id,sex}
  [{id: 1}, {id: 2}] NOT POSSIBLE FOR NON COMPOSITE
COMPOSITE PK:
  [1, 2]                                                                        -> {owner: [1], recipient: [2]}                                   {owner,recipient}
  [[1, 2], [3, 4]]                                                              -> {owner: [1, 2], recipient: [3, 4]}                             {owner,recipient}
  {owner: 1, recipient: 2, sex: 0}                                              -> {owner: [1], recipient: [2], sex: [0]}                         {owner,recipient,sex}
  [{owner: 1, recipient: 2}, {owner: 3, recipient: 4}]                          -> {owner: [1, 3], recipient: [2, 4]}                             {owner,recipient}
  [{owner: 1, recipient: 2, sex: 0}, {owner: 3, recipient: 4, sex: 1}] NOT POSSIBLE, MUST MATCH EXACTLY THE PK

  [{owner: [1], recipient: [2], sex: 0}                                         -> {owner: [1], recipient: [2], sex: [0]} // NOT A PK             {owner,recipient,sex}
  [{owner: [1], recipient: [2], sex: 0}, {owner: 3, recipient: 4, sex: 1}]      -> {owner: [1, 3], recipient: [2, 4], sex: [0, 1]} // NOT A PK    {owner,recipient,sex}
*/

const asc = (a: string, b: string): number => a.localeCompare(b);
const notNull = (el: string | undefined): boolean => el != null;

function getNewFiltersAndMapKeys<K extends object>(
  cur: FilterQueryDataloader<K>,
  meta: EntityMetadata<K>,
  entityName: string,
): Array<[FilterQueryDataloader<K>, string]>;
function getNewFiltersAndMapKeys<K extends object>(
  cur: FilterQueryDataloader<K>,
  meta: EntityMetadata<K>,
): [FilterQueryDataloader<K>, string];
function getNewFiltersAndMapKeys<K extends object>(
  cur: FilterQueryDataloader<K>,
  meta: EntityMetadata<K>,
  entityName?: string,
): [FilterQueryDataloader<K>, string] | Array<[FilterQueryDataloader<K>, string]> {
  const PKs = getPKs(cur, meta);
  if (PKs != null) {
    const res: [FilterQueryDataloader<K>, string] = [
      Object.fromEntries(
        meta.primaryKeys.map<[string, any]>((pk, i) => [pk, meta.compositePK ? PKs[i] : PKs]),
      ) as FilterQueryDataloader<K>,
      [entityName, `{${meta.primaryKeys.sort(asc).join(",")}}`].filter(notNull).join("|"),
    ];
    return entityName == null ? res : [res];
  } else {
    const newFilter: any = {};
    const keys: string[] = [];
    if (Array.isArray(cur)) {
      // COMPOSITE PKs like [{owner: 1, recipient: 2}, {recipient: 4, owner: 3}]
      for (const key of meta.primaryKeys) {
        newFilter[key] = cur.map((el: any) => {
          if (el[key] == null) {
            throw new Error(`Invalid query, missing composite PK ${key}`);
          }
          return el[key];
        });
        keys.push(key);
      }
      return [newFilter, `{${keys.sort(asc).join(",")}}`];
    } else {
      for (const [key, value] of Object.entries<any>(cur)) {
        // 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))
            .flat();
        }
        const keyProp = meta.properties[key as EntityKey<K>];
        if (keyProp == null) {
          throw new Error(`Cannot find properties for ${key}`);
        }
        if (keyProp.targetMeta == null) {
          newFilter[key] = Array.isArray(value) ? value : [value];
          keys.push(key);
        } else {
          const [subFilter, subKey] = getNewFiltersAndMapKeys(value, keyProp.targetMeta);
          newFilter[key] = subFilter;
          keys.push(`${key}:${subKey}`);
        }
      }
      const res: [FilterQueryDataloader<K>, string] = [
        newFilter,
        [entityName, `{${keys.sort(asc).join(",")}}`].filter(notNull).join("|"),
      ];
      return entityName == null ? res : [res];
    }
  }
}

function updateQueryFilter<K extends object, P extends string = never>(
  [acc, accOptions]: [FilterQueryDataloader<K>, { populate?: true | Set<any> }?],
  cur: FilterQueryDataloader<K>,
  options?: Pick<FindOptions<K, P>, "populate">,
): void {
  if (options?.populate != null && accOptions != null && accOptions.populate !== true) {
    if (Array.isArray(options.populate) && options.populate[0] === "*") {
      accOptions.populate = true;
    } else if (Array.isArray(options.populate)) {
      if (accOptions.populate == null) {
        accOptions.populate = new Set(options.populate);
      } else {
        for (const el of options.populate) {
          accOptions.populate.add(el);
        }
      }
    }
  }
  for (const [key, value] of Object.entries(acc)) {
    // 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), []));
    } else {
      updateQueryFilter([value], curValue);
    }
  }
}

export interface DataloaderFind<K extends object, Hint extends string = never, Fields extends string = never> {
  entityName: string;
  meta: EntityMetadata<K>;
  filter: FilterQueryDataloader<K>;
  options?: Pick<FindOptions<K, Hint, Fields>, "populate">;
  filtersAndKeys: Array<{ key: string; newFilter: FilterQueryDataloader<K> }>;
  many: boolean;
}

export function groupFindQueriesByOpts(
  dataloaderFinds: Array<PartialBy<DataloaderFind<any, any>, "filtersAndKeys">>,
): Map<string, [FilterQueryDataloader<any>, { populate?: true | Set<any> }?]> {
  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);
    dataloaderFind.filtersAndKeys = [];
    filtersAndKeys.forEach(([newFilter, key]) => {
      dataloaderFind.filtersAndKeys?.push({ key, newFilter });
      let queryMap = queriesMap.get(key);
      if (queryMap == null) {
        queryMap = [structuredClone(newFilter), {}];
        queriesMap.set(key, queryMap);
      } else {
        updateQueryFilter(queryMap, newFilter, options);
      }
    });
  }
  return queriesMap;
}

export function getFindBatchLoadFn<Entity extends object>(
  em: EntityManager,
  entityName: EntityName<Entity>,
): DataLoader.BatchLoadFn<PartialBy<DataloaderFind<any, any>, "filtersAndKeys">, any> {
  return async (dataloaderFinds: Array<PartialBy<DataloaderFind<any, any>, "filtersAndKeys">>) => {
    const optsMap = groupFindQueriesByOpts(dataloaderFinds);
    assertHasNewFilterAndMapKey(dataloaderFinds);

    const promises = optsMapToQueries(optsMap, em, entityName);
    const resultsMap = new Map(await Promise.all(promises));

    return dataloaderFinds.map(({ filtersAndKeys, many }) => {
      const res = filtersAndKeys.reduce<any[]>((acc, { key, newFilter }) => {
        const entities = resultsMap.get(key);
        if (entities == null) {
          // Should never happen
          /* istanbul ignore next */
          throw new Error("Cannot match results");
        }
        const res = entities[many ? "filter" : "find"]((entity) => {
          return filterResult(entity, newFilter);
        });
        acc.push(...(Array.isArray(res) ? res : [res]));
        return acc;
      }, []);
      return many ? res : res[0] ?? null;
    });

    function filterResult<K extends object>(entity: K, filter: FilterQueryDataloader<K>): boolean {
      for (const [key, value] of Object.entries(filter)) {
        const entityValue = entity[key as keyof K];
        if (Array.isArray(value)) {
          if (Array.isArray(entityValue)) {
            // Collection
            if (!value.every((el) => entityValue.includes(el))) {
              return false;
            }
          } else {
            // Single value
            if (!value.includes(entityValue)) {
              return false;
            }
          }
        } else {
          // Object: recursion
          if (!filterResult(entityValue as object, value)) {
            return false;
          }
        }
      }
      return true;
    }
  };
}

export function optsMapToQueries<Entity extends object>(
  optsMap: Map<string, [FilterQueryDataloader<any>, { populate?: true | Set<any> }?]>,
  em: EntityManager,
  entityName: EntityName<Entity>,
): Array<Promise<[string, any[]]>> {
  return Array.from(optsMap, async ([key, [filter, options]]): Promise<[string, any[]]> => {
    const findOptions = {
      ...(options?.populate != null && {
        populate: options.populate === true ? ["*"] : Array.from(options.populate),
      }),
    } satisfies Pick<FindOptions<any, any>, "populate">;
    const entities = await em.find(entityName, filter, findOptions);
    return [key, entities];
  });
}

export function assertHasNewFilterAndMapKey(
  dataloaderFinds: Array<PartialBy<DataloaderFind<any, any>, "filtersAndKeys">>,
): asserts dataloaderFinds is Array<DataloaderFind<any, any>> {
  /* if (dataloaderFinds.some((el) => el.key == null || el.newFilter == null)) {
    throw new Error("Missing key or newFilter");
  } */
}

export function hasRef<T extends AnyEntity<T>>(entity: T): T & Record<string, Reference<T>> {
  return entity as T & Record<string, Reference<T>>;
}

export function hasCol<T extends AnyEntity<T>, K extends object>(entity: T): T & Record<string, Collection<T, K>> {
  return entity as T & Record<string, Collection<T, K>>;
}

export function isRef<T extends AnyEntity<T>, K extends object>(
  refOrCol: Reference<T> | Collection<T, K>,
): refOrCol is Reference<T> {
  return !(refOrCol instanceof Collection);
}

export function isCol<T extends AnyEntity<T>, K extends object>(
  refOrCol: Reference<T> | Collection<T, K>,
): refOrCol is Collection<T, K> {
  return refOrCol instanceof Collection;
}