Add loadableGroup method to dataloader plugin

packages/plugin-dataloader/src/field-builder.ts

@@ -11,7 +11,7 @@ import {
   ShapeFromTypeParam,
   TypeParam,
 } from '@pothos/core';
-import type { LoadableListFieldOptions } from './types';
+import type { LoadableGroupFieldOptions, LoadableListFieldOptions } from './types';
 import { LoadableFieldOptions, LoaderShapeFromType } from './types';
 import { dataloaderGetter, rejectErrors } from './util';
 
@@ -137,3 +137,72 @@ fieldBuilderProto.loadableList = function loadableList<
     },
   });
 };
+
+fieldBuilderProto.loadableGroup = function loadableGroup<
+  Args extends InputFieldMap,
+  Type extends OutputType<SchemaTypes>,
+  Key,
+  CacheKey,
+  ResolveReturnShape,
+  Nullable extends FieldNullability<[Type]> = SchemaTypes['DefaultFieldNullability'],
+>({
+  load,
+  group,
+  loaderOptions,
+  resolve,
+  type,
+  ...options
+}: LoadableGroupFieldOptions<
+  SchemaTypes,
+  unknown,
+  Type,
+  Nullable,
+  Args,
+  ResolveReturnShape,
+  Key,
+  CacheKey,
+  FieldKind
+>): FieldRef<unknown> {
+  const getLoader = dataloaderGetter<Key, ShapeFromTypeParam<SchemaTypes, Type, true>[], CacheKey>(
+    loaderOptions,
+    async (keys, ctx) => {
+      const values = await load(keys, ctx);
+      const groups = new Map<Key, ShapeFromTypeParam<SchemaTypes, Type, true>[]>();
+
+      for (const value of values) {
+        if (value == null) {
+          // eslint-disable-next-line no-continue
+          continue;
+        }
+
+        const groupKey = group(value);
+        if (!groups.has(groupKey)) {
+          groups.set(groupKey, []);
+        }
+
+        groups.get(groupKey)!.push(value);
+      }
+
+      return keys.map((key) => groups.get(key) ?? []);
+    },
+    undefined,
+    false,
+  );
+
+  return this.field({
+    ...options,
+    type: [type],
+    // @ts-expect-error types don't match because this handles both lists and single objects
+    resolve: async (
+      parent: unknown,
+      args: InputShapeFromFields<Args>,
+      context: {},
+      info: GraphQLResolveInfo,
+    ) => {
+      const ids = await resolve(parent, args, context, info);
+      const loader = getLoader(context);
+
+      return loader.load(ids as Key);
+    },
+  });
+};

packages/plugin-dataloader/src/global-types.ts

@@ -24,6 +24,7 @@ import {
 } from './types';
 
 import type {
+  LoadableGroupFieldOptions,
   LoadableInterfaceOptions,
   LoadableListFieldOptions,
   LoadableUnionOptions,
@@ -175,6 +176,27 @@ declare global {
           Kind
         >,
       ) => FieldRef<unknown>;
+
+      loadableGroup: <
+        Args extends InputFieldMap,
+        Type extends OutputType<Types>,
+        Key,
+        CacheKey,
+        ResolveReturnShape,
+        Nullable extends FieldNullability<[Type]> = Types['DefaultFieldNullability'],
+      >(
+        options: LoadableGroupFieldOptions<
+          Types,
+          ParentShape,
+          Type,
+          Nullable,
+          Args,
+          ResolveReturnShape,
+          Key,
+          CacheKey,
+          Kind
+        >,
+      ) => FieldRef<unknown>;
     }
   }
 }

packages/plugin-dataloader/src/types.ts

@@ -90,6 +90,36 @@ export type LoadableListFieldOptions<
   >;
 };
 
+export type LoadableGroupFieldOptions<
+  Types extends SchemaTypes,
+  ParentShape,
+  Type extends OutputType<Types>,
+  Nullable extends FieldNullability<[Type]>,
+  Args extends InputFieldMap,
+  ResolveReturnShape,
+  Key,
+  CacheKey,
+  Kind extends FieldKind = FieldKind,
+> = Omit<
+  FieldOptionsFromKind<Types, ParentShape, [Type], Nullable, Args, Kind, Key, ResolveReturnShape>,
+  'resolve' | 'type'
+> & {
+  type: Type;
+  load: (
+    keys: Key[],
+    context: Types['Context'],
+  ) => Promise<readonly ShapeFromTypeParam<Types, Type, true>[]>;
+  loaderOptions?: DataLoader.Options<Key, ShapeFromTypeParam<Types, Type, true>[], CacheKey>;
+  group: (value: ShapeFromTypeParam<Types, Type, false>) => Key;
+  resolve: Resolver<
+    ParentShape,
+    InputShapeFromFields<Args>,
+    Types['Context'],
+    Key,
+    ResolveReturnShape
+  >;
+};
+
 export type DataLoaderOptions<
   Types extends SchemaTypes,
   Shape,

website/pages/docs/plugins/dataloader.mdx

@@ -160,6 +160,26 @@ builder.objectField(User, 'posts', (t) =>
 );
 ```
 
+### loadableGroup fields for one-to-many relations
+
+In many cases, it's easier to load a flat list in a dataloader rather than loading a list of lists.
+the `loadableGroup` method simplifies this.
+
+```typescript
+// Loading multiple Posts
+builder.objectField(User, 'posts', (t) =>
+  t.loadableGroup({
+    // type is singular, but will create a list field
+    type: Post,
+    // will be called with ids of all the users, and should return `Post[]`
+    load: (ids: number[], context) => db.posts.findMany({ where: { authorId: { in: ids } } }),
+    // will be called with each post to determine which group it belongs to
+    group: (post) => post.authorId,
+    resolve: (user, args) => user.id,
+  }),
+);
+```
+
 ### dataloader options
 
 You can provide additional options for your dataloaders using `loaderOptions`.