DataLoader Pattern

npm install dataloader

import DataLoader from 'dataloader';

// The batch function receives ALL keys collected in one tick
const userLoader = new DataLoader<string, User>(async (keys) => {
  console.log('Batch loading users:', keys);
  // keys might be: ['user_1', 'user_3', 'user_5']

  // One database query for all keys
  const users = await db.query('SELECT * FROM users WHERE id = ANY($1)', [keys]);

  // CRITICAL: Return results in the SAME ORDER as keys
  const userMap = new Map(users.map(u => [u.id, u]));
  return keys.map(key => userMap.get(key) || new Error(`User ${key} not found`));
});

// These three calls happen in the same tick
const promise1 = userLoader.load('user_1');  // Queued
const promise2 = userLoader.load('user_3');  // Queued
const promise3 = userLoader.load('user_5');  // Queued

// On the next tick, DataLoader calls the batch function once:
// batchFn(['user_1', 'user_3', 'user_5'])

const [user1, user3, user5] = await Promise.all([promise1, promise2, promise3]);

const userLoader = new DataLoader<string, User>(batchUsers);

// First load — hits the database
const user1a = await userLoader.load('user_1');

// Second load — returns cached result, no database call
const user1b = await userLoader.load('user_1');

console.log(user1a === user1b); // true — same object reference

// Clear cache for a specific key (after a mutation)
userLoader.clear('user_1');

// Clear all cache
userLoader.clearAll();

// Loading posts for multiple authors in one query
const postsByAuthorLoader = new DataLoader<string, Post[]>(
  async (authorIds) => {
    const posts = await prisma.post.findMany({
      where: { authorId: { in: [...authorIds] } },
      orderBy: { createdAt: 'desc' },
    });

    // Group by authorId and maintain input order
    const grouped = new Map<string, Post[]>();
    authorIds.forEach(id => grouped.set(id, []));
    posts.forEach(post => grouped.get(post.authorId)?.push(post));

    return authorIds.map(id => grouped.get(id) || []);
  }
);

// Loading comments count for multiple posts
const commentsCountLoader = new DataLoader<string, number>(
  async (postIds) => {
    const counts = await prisma.comment.groupBy({
      by: ['postId'],
      where: { postId: { in: [...postIds] } },
      _count: true,
    });

    const countMap = new Map(counts.map(c => [c.postId, c._count]));
    return postIds.map(id => countMap.get(id) || 0);
  }
);

// When you need to batch on multiple parameters
interface LikeKey {
  userId: string;
  postId: string;
}

const isLikedLoader = new DataLoader<LikeKey, boolean>(
  async (keys) => {
    const likes = await prisma.like.findMany({
      where: {
        OR: keys.map(k => ({
          userId: k.userId,
          postId: k.postId,
        })),
      },
    });

    const likeSet = new Set(
      likes.map(l => `${l.userId}:${l.postId}`)
    );

    return keys.map(k => likeSet.has(`${k.userId}:${k.postId}`));
  },
  {
    // Custom cache key function for composite keys
    cacheKeyFn: (key) => `${key.userId}:${key.postId}`,
  }
);

// Usage in resolver
const resolvers = {
  Post: {
    isLikedByViewer: (post: Post, _: unknown, ctx: Context) => {
      if (!ctx.currentUser) return false;
      return ctx.loaders.isLikedLoader.load({
        userId: ctx.currentUser.id,
        postId: post.id,
      });
    },
  },
};

import DataLoader from 'dataloader';
import { PrismaClient, User, Post, Comment } from '@prisma/client';

interface Loaders {
  userLoader: DataLoader<string, User>;
  postsByAuthorLoader: DataLoader<string, Post[]>;
  commentsByPostLoader: DataLoader<string, Comment[]>;
  commentsCountLoader: DataLoader<string, number>;
  isLikedLoader: DataLoader<{ userId: string; postId: string }, boolean>;
}

export interface Context {
  prisma: PrismaClient;
  currentUser: User | null;
  loaders: Loaders;
}

export function createLoaders(prisma: PrismaClient): Loaders {
  return {
    userLoader: new DataLoader(async (ids: readonly string[]) => {
      const users = await prisma.user.findMany({
        where: { id: { in: [...ids] } },
      });
      const map = new Map(users.map(u => [u.id, u]));
      return ids.map(id => map.get(id)!);
    }),

    postsByAuthorLoader: new DataLoader(async (authorIds: readonly string[]) => {
      const posts = await prisma.post.findMany({
        where: { authorId: { in: [...authorIds] } },
      });
      const grouped = new Map<string, Post[]>();
      authorIds.forEach(id => grouped.set(id, []));
      posts.forEach(p => grouped.get(p.authorId)?.push(p));
      return authorIds.map(id => grouped.get(id) || []);
    }),

    commentsByPostLoader: new DataLoader(async (postIds: readonly string[]) => {
      const comments = await prisma.comment.findMany({
        where: { postId: { in: [...postIds] } },
      });
      const grouped = new Map<string, Comment[]>();
      postIds.forEach(id => grouped.set(id, []));
      comments.forEach(c => grouped.get(c.postId)?.push(c));
      return postIds.map(id => grouped.get(id) || []);
    }),

    commentsCountLoader: new DataLoader(async (postIds: readonly string[]) => {
      const counts = await prisma.comment.groupBy({
        by: ['postId'],
        where: { postId: { in: [...postIds] } },
        _count: true,
      });
      const map = new Map(counts.map(c => [c.postId, c._count]));
      return postIds.map(id => map.get(id) || 0);
    }),

    isLikedLoader: new DataLoader(
      async (keys: readonly { userId: string; postId: string }[]) => {
        const likes = await prisma.like.findMany({
          where: { OR: keys.map(k => ({ userId: k.userId, postId: k.postId })) },
        });
        const set = new Set(likes.map(l => `${l.userId}:${l.postId}`));
        return keys.map(k => set.has(`${k.userId}:${k.postId}`));
      },
      { cacheKeyFn: (key) => `${key.userId}:${key.postId}` }
    ),
  };
}