Merge pull request #15 from facebook/cachekey

Allow for custom cache key function

Author: leebyron

README.md

@@ -151,6 +151,10 @@ Create a new `DataLoader` given a batch loading function and options.
   - *cache*: Default `true`. Set to `false` to disable caching, instead
     creating a new Promise and new key in the `batchLoadFn` for every load.
 
+  - *cacheKeyFn*: A function to produce a cache key for a given load key.
+    Defaults to `key => key`. Useful to provide when JavaScript objects are keys
+    and two similarly shaped objects should be considered equivalent.
+
 ##### `load(key)`
 
 Loads a key, returning a `Promise` for the value represented by that key.

src/__tests__/dataloader-test.js

@@ -418,6 +418,56 @@ describe('Accepts options', () => {
     );
   });
 
+  describe('Accepts object key in custom cacheKey function', () => {
+    function cacheKey(key) {
+      var result;
+      if (typeof key === 'object') {
+        result = Object.keys(key).sort().map(k => k + ':' + key[k]).join('-');
+      } else {
+        result = String(key);
+      }
+      return result;
+    }
+
+    it('Accepts objects with simple key', async () => {
+      var keyA = '1234';
+      var identityLoadCalls = [];
+      var identityLoader = new DataLoader(keys => {
+        identityLoadCalls.push(keys);
+        return Promise.resolve(keys);
+      }, { cacheKeyFn: cacheKey });
+
+      var valueA = await identityLoader.load(keyA);
+      expect(valueA).to.equal(keyA);
+    });
+
+    it('Accepts objects with different order of keys', async () => {
+      var keyA = { a: 123, b: 321 };
+      var keyB = { b: 321, a: 123 };
+
+      var identityLoadCalls = [];
+      var identityLoader = new DataLoader(keys => {
+        identityLoadCalls.push(keys);
+        return Promise.resolve(keys);
+      }, { cacheKeyFn: cacheKey });
+
+      // Fetches as expected
+
+      var [ valueA, valueB ] = await Promise.all([
+        identityLoader.load(keyA),
+        identityLoader.load(keyB),
+      ]);
+
+      expect(valueA).to.equal(keyA);
+      expect(valueB).to.equal(valueA);
+
+      expect(identityLoadCalls).to.have.length(1);
+      expect(identityLoadCalls[0]).to.have.length(1);
+      expect(identityLoadCalls[0][0]).to.equal(keyA);
+    });
+
+  });
+
 });
 
 describe('It is resilient to job queue ordering', () => {

src/index.js

@@ -12,8 +12,12 @@
 // of values or Errors.
 type BatchLoadFn<K, V> = (keys: Array<K>) => Promise<Array<V | Error>>
 
-// Optionally turn off batching or caching.
-type Options = { batch?: boolean, cache?: boolean }
+// Optionally turn off batching or caching or provide a cache key function.
+type Options = {
+  batch?: boolean,
+  cache?: boolean,
+  cacheKeyFn?: (key: any) => any
+}
 
 /**
  * A `DataLoader` creates a public API for loading data from a particular
@@ -63,10 +67,12 @@ export default 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;
 
     // If caching and there is a cache-hit, return cached Promise.
     if (shouldCache) {
-      var cachedPromise = this._promiseCache.get(key);
+      var cachedPromise = this._promiseCache.get(cacheKey);
       if (cachedPromise) {
         return cachedPromise;
       }
@@ -93,7 +99,7 @@ export default class DataLoader<K, V> {
 
     // If caching, cache this promise.
     if (shouldCache) {
-      this._promiseCache.set(key, promise);
+      this._promiseCache.set(cacheKey, promise);
     }
 
     return promise;
@@ -127,7 +133,9 @@ export default class DataLoader<K, V> {
    * method chaining.
    */
   clear(key: K): DataLoader<K, V> {
-    this._promiseCache.delete(key);
+    var cacheKeyFn = this._options && this._options.cacheKeyFn;
+    var cacheKey = cacheKeyFn ? cacheKeyFn(key) : key;
+    this._promiseCache.delete(cacheKey);
     return this;
   }