Allow specific CacheMap key type (#145)

* Allow specific CacheMap key type

* Update impl, flow types, and tests

Author: AdamHerrmann

src/__tests__/dataloader.test.js

@@ -7,9 +7,12 @@
  * @flow
  */
 
+import type { Options } from '..';
 const DataLoader = require('..');
 
-function idLoader(options) {
+function idLoader<K, C>(
+  options?: Options<K, K, C>
+): [ DataLoader<K, K, C>, Array<$ReadOnlyArray<K>> ] {
   const loadCalls = [];
   const identityLoader = new DataLoader(keys => {
     loadCalls.push(keys);
@@ -568,7 +571,7 @@ describe('Accepts options', () => {
   });
 
   describe('Accepts object key in custom cacheKey function', () => {
-    function cacheKey(key) {
+    function cacheKey(key: {[string]: any}): string {
       return Object.keys(key).sort().map(k => k + ':' + key[k]).join();
     }
 

src/index.d.ts

@@ -15,9 +15,9 @@
  * with different access permissions and consider creating a new instance
  * per web request.
  */
-declare class DataLoader<K, V> {
+declare class DataLoader<K, V, C = K> {
 
-  constructor(batchLoadFn: DataLoader.BatchLoadFn<K, V>, options?: DataLoader.Options<K, V>);
+  constructor(batchLoadFn: DataLoader.BatchLoadFn<K, V>, options?: DataLoader.Options<K, V, C>);
 
   /**
    * Loads a key, returning a `Promise` for the value represented by that key.
@@ -43,20 +43,20 @@ declare class DataLoader<K, V> {
    * Clears the value at `key` from the cache, if it exists. Returns itself for
    * method chaining.
    */
-  clear(key: K): DataLoader<K, V>;
+  clear(key: K): this;
 
   /**
    * Clears the entire cache. To be used when some event results in unknown
    * invalidations across this particular `DataLoader`. Returns itself for
    * method chaining.
    */
-  clearAll(): DataLoader<K, V>;
+  clearAll(): this;
 
   /**
    * Adds the provied key and value to the cache. If the key already exists, no
    * change is made. Returns itself for method chaining.
    */
-  prime(key: K, value: V): DataLoader<K, V>;
+  prime(key: K, value: V): this;
 }
 
 declare namespace DataLoader {
@@ -74,7 +74,7 @@ declare namespace DataLoader {
 
   // Optionally turn off batching or caching or provide a cache key function or a
   // custom cache instance.
-  export type Options<K, V> = {
+  export type Options<K, V, C = K> = {
 
     /**
      * Default `true`. Set to `false` to disable batching,
@@ -102,14 +102,14 @@ declare namespace DataLoader {
      * objects are keys and two similarly shaped objects should
      * be considered equivalent.
      */
-    cacheKeyFn?: (key: any) => any,
+    cacheKeyFn?: (key: K) => C,
 
     /**
      * An instance of Map (or an object with a similar API) to
      * be used as the underlying cache for this loader.
      * Default `new Map()`.
      */
-    cacheMap?: CacheMap<K, Promise<V>>;
+    cacheMap?: CacheMap<C, Promise<V>>;
   }
 }
 

src/index.js

@@ -14,12 +14,12 @@ export type BatchLoadFn<K, V> =
 
 // Optionally turn off batching or caching or provide a cache key function or a
 // custom cache instance.
-export type Options<K, V> = {
+export type Options<K, V, C = K> = {
   batch?: boolean;
   maxBatchSize?: number;
   cache?: boolean;
-  cacheKeyFn?: (key: any) => any;
-  cacheMap?: CacheMap<K, Promise<V>>;
+  cacheKeyFn?: (key: K) => C;
+  cacheMap?: CacheMap<C, Promise<V>>;
 };
 
 // If a custom cache is provided, it must be of this type (a subset of ES6 Map).
@@ -40,10 +40,10 @@ export type CacheMap<K, V> = {
  * different access permissions and consider creating a new instance per
  * web request.
  */
-class DataLoader<K, V> {
+class DataLoader<K, V, C = K> {
   constructor(
     batchLoadFn: BatchLoadFn<K, V>,
-    options?: Options<K, V>
+    options?: Options<K, V, C>
   ) {
     if (typeof batchLoadFn !== 'function') {
       throw new TypeError(
@@ -59,8 +59,8 @@ class DataLoader<K, V> {
 
   // Private
   _batchLoadFn: BatchLoadFn<K, V>;
-  _options: ?Options<K, V>;
-  _promiseCache: CacheMap<K, Promise<V>>;
+  _options: ?Options<K, V, C>;
+  _promiseCache: CacheMap<C, Promise<V>>;
   _queue: LoaderQueue<K, V>;
 
   /**
@@ -78,8 +78,7 @@ class DataLoader<K, V> {
     var options = this._options;
     var shouldBatch = !options || options.batch !== false;
     var shouldCache = !options || options.cache !== false;
-    var cacheKeyFn = options && options.cacheKeyFn;
-    var cacheKey = cacheKeyFn ? cacheKeyFn(key) : key;
+    var cacheKey = getCacheKey(options, key);
 
     // If caching and there is a cache-hit, return cached Promise.
     if (shouldCache) {
@@ -143,9 +142,8 @@ class DataLoader<K, V> {
    * Clears the value at `key` from the cache, if it exists. Returns itself for
    * method chaining.
    */
-  clear(key: K): DataLoader<K, V> {
-    var cacheKeyFn = this._options && this._options.cacheKeyFn;
-    var cacheKey = cacheKeyFn ? cacheKeyFn(key) : key;
+  clear(key: K): this {
+    var cacheKey = getCacheKey(this._options, key);
     this._promiseCache.delete(cacheKey);
     return this;
   }
@@ -155,7 +153,7 @@ class DataLoader<K, V> {
    * invalidations across this particular `DataLoader`. Returns itself for
    * method chaining.
    */
-  clearAll(): DataLoader<K, V> {
+  clearAll(): this {
     this._promiseCache.clear();
     return this;
   }
@@ -164,9 +162,8 @@ class DataLoader<K, V> {
    * Adds the provided key and value to the cache. If the key already
    * exists, no change is made. Returns itself for method chaining.
    */
-  prime(key: K, value: V): DataLoader<K, V> {
-    var cacheKeyFn = this._options && this._options.cacheKeyFn;
-    var cacheKey = cacheKeyFn ? cacheKeyFn(key) : key;
+  prime(key: K, value: V): this {
+    var cacheKey = getCacheKey(this._options, key);
 
     // Only add the key if it does not already exist.
     if (this._promiseCache.get(cacheKey) === undefined) {
@@ -224,7 +221,7 @@ var resolvedPromise;
 
 // Private: given the current state of a Loader instance, perform a batch load
 // from its current queue.
-function dispatchQueue<K, V>(loader: DataLoader<K, V>) {
+function dispatchQueue<K, V>(loader: DataLoader<K, V, any>) {
   // Take the current loader queue, replacing it with an empty queue.
   var queue = loader._queue;
   loader._queue = [];
@@ -245,7 +242,7 @@ function dispatchQueue<K, V>(loader: DataLoader<K, V>) {
 }
 
 function dispatchQueueBatch<K, V>(
-  loader: DataLoader<K, V>,
+  loader: DataLoader<K, V, any>,
   queue: LoaderQueue<K, V>
 ) {
   // Collect all keys to be loaded in this dispatch
@@ -303,7 +300,7 @@ function dispatchQueueBatch<K, V>(
 // Private: do not cache individual loads if the entire batch dispatch fails,
 // but still reject each request so they do not hang.
 function failedDispatch<K, V>(
-  loader: DataLoader<K, V>,
+  loader: DataLoader<K, V, any>,
   queue: LoaderQueue<K, V>,
   error: Error
 ) {
@@ -313,10 +310,19 @@ function failedDispatch<K, V>(
   });
 }
 
+// Private: produce a cache key for a given key (and options)
+function getCacheKey<K, V, C>(
+  options: ?Options<K, V, C>,
+  key: K
+): C {
+  var cacheKeyFn = options && options.cacheKeyFn;
+  return cacheKeyFn ? cacheKeyFn(key) : (key: any);
+}
+
 // Private: given the DataLoader's options, produce a CacheMap to be used.
-function getValidCacheMap<K, V>(
-  options: ?Options<K, V>
-): CacheMap<K, Promise<V>> {
+function getValidCacheMap<K, V, C>(
+  options: ?Options<K, V, C>
+): CacheMap<C, Promise<V>> {
   var cacheMap = options && options.cacheMap;
   if (!cacheMap) {
     return new Map();