|  | 
| 1 |  | -import type DataLoader from "dataloader"; | 
| 2 |  | -import type Keyv from "keyv"; | 
| 3 |  | -import hashObject from "object-hash"; | 
| 4 |  | - | 
| 5 |  | -type NotUndefined = object | string | number | boolean | NotUndefined[]; | 
| 6 |  | - | 
| 7 |  | -export type cacheOptions<K, V> = { | 
| 8 |  | -	keys: ReadonlyArray<K>; | 
| 9 |  | -	store?: Keyv<V>; | 
| 10 |  | -	ttl: number; | 
| 11 |  | - | 
| 12 |  | -	batchLoadFn: DataLoader.BatchLoadFn<K, V>; | 
| 13 |  | -	cacheKeysFn: (ref: K) => string[]; | 
| 14 |  | -}; | 
| 15 |  | - | 
| 16 |  | -// dataloaderCache is a wrapper around the dataloader batchLoadFn that adds | 
| 17 |  | -// caching. It takes an array of keys and returns an array of items. If an item | 
| 18 |  | -// is not in the cache, it will be fetched from the batchLoadFn and then written | 
| 19 |  | -// to the cache for next time. | 
| 20 |  | -// | 
| 21 |  | -// Note: this function is O^2, so it should only be used for small batches of | 
| 22 |  | -// keys. | 
| 23 |  | -export const dataloaderCache = async <K extends NotUndefined, V>( | 
| 24 |  | -	args: cacheOptions<K, V | null>, | 
| 25 |  | -): Promise<(V | null)[]> => { | 
| 26 |  | -	const result = await fromCache<K, V>(args.keys, args); | 
| 27 |  | -	const store: Record<string, V | null> = {}; | 
| 28 |  | - | 
| 29 |  | -	// Check results, if an item is null then it was not in the cache, we place | 
| 30 |  | -	// these in the cacheMiss array and fetch them. | 
| 31 |  | -	const cacheMiss: Array<K> = []; | 
| 32 |  | -	for (const [key, cached] of zip(args.keys, result)) { | 
| 33 |  | -		if (cached === undefined) { | 
| 34 |  | -			cacheMiss.push(key); | 
| 35 |  | -		} else { | 
| 36 |  | -			store[hashObject(key)] = cached; | 
| 37 |  | -		} | 
| 38 |  | -	} | 
| 39 |  | - | 
| 40 |  | -	// Fetch the items that are not in the cache and write them to the cache for | 
| 41 |  | -	// next time | 
| 42 |  | -	if (cacheMiss.length > 0) { | 
| 43 |  | -		const newItems = await args.batchLoadFn(cacheMiss); | 
| 44 |  | -		const buffer = new Map<string, V | null>(); | 
| 45 |  | - | 
| 46 |  | -		for (const [key, item] of zip(cacheMiss, Array.from(newItems))) { | 
| 47 |  | -			if (key === undefined) { | 
| 48 |  | -				throw new Error("key is undefined"); | 
| 49 |  | -			} | 
| 50 |  | - | 
| 51 |  | -			if (!(item instanceof Error)) { | 
| 52 |  | -				store[hashObject(key)] = item; | 
| 53 |  | - | 
| 54 |  | -				const cacheKeys = args.cacheKeysFn(key); | 
| 55 |  | -				for (const cacheKey of cacheKeys) { | 
| 56 |  | -					buffer.set(cacheKey, item); | 
| 57 |  | -				} | 
| 58 |  | -			} | 
| 59 |  | -		} | 
| 60 |  | - | 
| 61 |  | -		await toCache<K, V | null>(buffer, args); | 
| 62 |  | -	} | 
| 63 |  | - | 
| 64 |  | -	return args.keys.map((key) => { | 
| 65 |  | -		const item = store[hashObject(key)]; | 
| 66 |  | -		if (item) { | 
| 67 |  | -			return item; | 
| 68 |  | -		} | 
| 69 |  | - | 
| 70 |  | -		return null; | 
| 71 |  | -	}); | 
| 72 |  | -}; | 
| 73 |  | - | 
| 74 |  | -// Read items from the cache by the keys | 
| 75 |  | -const fromCache = async <K, V>( | 
| 76 |  | -	keys: ReadonlyArray<K>, | 
| 77 |  | -	options: cacheOptions<K, V | null>, | 
| 78 |  | -): Promise<(V | null | undefined)[]> => { | 
| 79 |  | -	if (!options.store) { | 
| 80 |  | -		return new Array<V | null | undefined>(keys.length).fill(undefined); | 
| 81 |  | -	} | 
| 82 |  | - | 
| 83 |  | -	const cacheKeys = keys.flatMap(options.cacheKeysFn); | 
| 84 |  | -	const cachedValues = await options.store.get(cacheKeys); | 
| 85 |  | -	return cachedValues.map((v) => v); | 
| 86 |  | -}; | 
| 87 |  | - | 
| 88 |  | -// Write items to the cache | 
| 89 |  | -const toCache = async <K, V>( | 
| 90 |  | -	items: Map<string, V | null>, | 
| 91 |  | -	options: cacheOptions<K, V | null>, | 
| 92 |  | -): Promise<void> => { | 
| 93 |  | -	if (!options.store) { | 
| 94 |  | -		return; | 
| 95 |  | -	} | 
| 96 |  | -	for (const [key, value] of items) { | 
| 97 |  | -		options.store.set(key, value, options.ttl); | 
| 98 |  | -	} | 
| 99 |  | -}; | 
| 100 |  | - | 
| 101 |  | -function zip<T, U>(arr1: readonly T[], arr2: readonly U[]): [T, U][] { | 
| 102 |  | -	const minLength = Math.min(arr1.length, arr2.length); | 
| 103 |  | -	const result: [T, U][] = []; | 
| 104 |  | - | 
| 105 |  | -	for (let i = 0; i < minLength; i++) { | 
| 106 |  | -		result.push([arr1[i], arr2[i]]); | 
| 107 |  | -	} | 
| 108 |  | - | 
| 109 |  | -	return result; | 
| 110 |  | -} | 
|  | 1 | +export { dataloaderCache } from "./cache"; | 
|  | 2 | +export type { cacheOptions } from "./cache"; | 
|  | 3 | +export { DataLoaderCache } from "./wrapper"; | 
0 commit comments