DataLoader

const DataLoader = require('dataloader');

import DataLoader from 'dataloader';
// Or for CommonJS export compatibility:
import DataLoader = require('dataloader');

const DataLoader = require('dataloader');

// Create a batch loading function
async function batchGetUsers(userIds) {
  const users = await db.query('SELECT * FROM users WHERE id IN (?)', [userIds]);
  // Return results in same order as input keys
  return userIds.map(id => users.find(user => user.id === id) || new Error(`User ${id} not found`));
}

// Create a DataLoader instance
const userLoader = new DataLoader(batchGetUsers);

// Load individual users - these will be batched automatically
const user1 = await userLoader.load(1);
const user2 = await userLoader.load(2);

// Or load multiple users at once
const [user3, user4] = await userLoader.loadMany([3, 4]);

/**
 * Creates a new DataLoader instance
 * @param batchLoadFn - Function that loads multiple keys and returns Promise of results
 * @param options - Optional configuration object
 */
constructor<K, V, C = K>(
  batchLoadFn: BatchLoadFn<K, V>,
  options?: Options<K, V, C>
): DataLoader<K, V, C>;

type BatchLoadFn<K, V> = (keys: ReadonlyArray<K>) => Promise<ArrayLike<V | Error>>;

interface Options<K, V, C = K> {
  /** Enable/disable batching (default: true) */
  batch?: boolean;
  /** Maximum keys per batch (default: Infinity) */
  maxBatchSize?: number;
  /** Custom batch scheduling function */
  batchScheduleFn?: (callback: () => void) => void;
  /** Enable/disable caching (default: true) */
  cache?: boolean;
  /** Function to generate cache keys from load keys */
  cacheKeyFn?: (key: K) => C;
  /** Custom cache implementation (default: new Map()) */
  cacheMap?: CacheMap<C, Promise<V>> | null;
  /** Name for this DataLoader instance (useful for APM tools) */
  name?: string | null;
}

interface CacheMap<K, V> {
  get(key: K): V | void;
  set(key: K, value: V): any;
  delete(key: K): any;
  clear(): any;
}

// Basic usage
const userLoader = new DataLoader(async (userIds) => {
  const users = await getUsersByIds(userIds);
  return userIds.map(id => users[id] || new Error(`User ${id} not found`));
});

// With options
const userLoader = new DataLoader(batchLoadFn, {
  batch: true,
  maxBatchSize: 100,
  cache: true,
  cacheKeyFn: key => key.toString(),
  name: 'UserLoader'
});

// Disable batching (equivalent to maxBatchSize: 1)
const immediateLoader = new DataLoader(batchLoadFn, { batch: false });

// Custom cache
const userLoader = new DataLoader(batchLoadFn, {
  cacheMap: new Map() // or null to disable caching
});

/**
 * Loads a key, returning a Promise for the value represented by that key
 * @param key - The key to load
 * @returns Promise that resolves to the value for the given key
 * @throws TypeError if key is null or undefined
 */
load(key: K): Promise<V>;

// Single load
const user = await userLoader.load(123);

// Multiple concurrent loads (automatically batched)
const [user1, user2, user3] = await Promise.all([
  userLoader.load(1),
  userLoader.load(2),
  userLoader.load(3)
]);

// Error handling
try {
  const user = await userLoader.load(999);
} catch (error) {
  console.log('User not found:', error.message);
}

/**
 * Loads multiple keys, promising an array of values or Error instances
 * @param keys - Array-like object containing keys to load
 * @returns Promise that resolves to array of values or Error instances
 * @throws TypeError if keys is not array-like
 */
loadMany(keys: ArrayLike<K>): Promise<Array<V | Error>>;

// Load multiple users
const results = await userLoader.loadMany([1, 2, 3, 999]);
// results might be: [User{id:1}, User{id:2}, User{id:3}, Error('User 999 not found')]

// Handle mixed results
const [user1, user2, user3, error] = await userLoader.loadMany([1, 2, 3, 999]);
if (error instanceof Error) {
  console.log('Failed to load user 999:', error.message);
}

// Filter out errors
const users = results.filter(result => !(result instanceof Error));

/**
 * Clears the value at key from the cache, if it exists
 * @param key - The key to clear from cache
 * @returns This DataLoader instance for method chaining
 */
clear(key: K): this;

// Clear a specific user from cache
userLoader.clear(123);

// Method chaining
userLoader
  .clear(1)
  .clear(2)
  .clear(3);

// Clear after data mutation
async function updateUser(id, updates) {
  const user = await db.updateUser(id, updates);
  userLoader.clear(id); // Remove stale cached data
  return user;
}

/**
 * Clears the entire cache
 * @returns This DataLoader instance for method chaining
 */
clearAll(): this;

// Clear all cached data
userLoader.clearAll();

// Clear cache after bulk operations
async function bulkUpdateUsers(updates) {
  await db.bulkUpdate(updates);
  userLoader.clearAll(); // Invalidate all cached users
}

// Periodic cache clearing
setInterval(() => {
  userLoader.clearAll();
}, 300000); // Clear cache every 5 minutes

/**
 * Adds the provided key and value to the cache
 * @param key - The key to prime
 * @param value - The value, Promise, or Error to associate with the key
 * @returns This DataLoader instance for method chaining
 */
prime(key: K, value: V | PromiseLike<V> | Error): this;

// Prime cache with known data
const newUser = await createUser({ name: 'Alice' });
userLoader.prime(newUser.id, newUser);

// Prime with error for known invalid keys
userLoader.prime(0, new Error('Invalid user ID'));

// Prime with promise
const userPromise = db.findUser(123);
userLoader.prime(123, userPromise);

// Method chaining
userLoader
  .prime(1, user1)
  .prime(2, user2)
  .prime(3, user3);

/**
 * The name given to this DataLoader instance
 * Useful for APM tools and debugging
 */
name: string | null;

// Set name in constructor
const userLoader = new DataLoader(batchLoadFn, { name: 'UserLoader' });
console.log(userLoader.name); // 'UserLoader'

// Access for debugging
function debugLoader(loader) {
  console.log(`DataLoader ${loader.name || 'unnamed'} cache status`);
}

// Individual errors in batch results
async function batchLoadUsers(ids) {
  const users = await db.getUsers(ids);
  return ids.map(id => {
    const user = users.find(u => u.id === id);
    return user || new Error(`User ${id} not found`);
  });
}

// Synchronous errors in batch function
const loader = new DataLoader((keys) => {
  if (keys.length > 100) {
    throw new Error('Too many keys requested');
  }
  return batchLoad(keys);
});

try {
  const user = await userLoader.load(invalidId);
} catch (error) {
  // Handle individual load errors
  console.error('Failed to load user:', error);
}

// With loadMany, errors are returned as Error instances
const results = await userLoader.loadMany([1, 2, invalidId]);
results.forEach((result, index) => {
  if (result instanceof Error) {
    console.error(`Failed to load key ${keys[index]}:`, result.message);
  }
});

// SQL with proper ordering
const userLoader = new DataLoader(async (userIds) => {
  const query = 'SELECT * FROM users WHERE id IN (?)';
  const users = await db.query(query, [userIds]);
  
  // Return results in same order as input keys
  return userIds.map(id => 
    users.find(user => user.id === id) || new Error(`User ${id} not found`)
  );
});

// Using with GraphQL resolvers
const resolvers = {
  Post: {
    author: (post) => userLoader.load(post.authorId),
  },
  Comment: {
    author: (comment) => userLoader.load(comment.authorId),
  }
};

// Express middleware
app.use((req, res, next) => {
  req.loaders = {
    user: new DataLoader(batchLoadUsers),
    post: new DataLoader(batchLoadPosts)
  };
  next();
});

// Use in routes
app.get('/posts/:id', async (req, res) => {
  const post = await req.loaders.post.load(req.params.id);
  const author = await req.loaders.user.load(post.authorId);
  res.json({ ...post, author });
});

// Case-insensitive user lookup
const userByEmailLoader = new DataLoader(
  async (emails) => {
    const users = await db.getUsersByEmail(emails);
    return emails.map(email => users[email.toLowerCase()] || null);
  },
  {
    cacheKeyFn: email => email.toLowerCase()
  }
);