Skip to content

Commit

Permalink
Allow specific CacheMap key type (#145)
Browse files Browse the repository at this point in the history 
* Allow specific CacheMap key type

* Update impl, flow types, and tests
  • Loading branch information
@AdamHerrmann @leebyron
AdamHerrmann authored and leebyron committed Nov 13, 2019
1 parent cd213db commit bf23eb9
Show file tree
Hide file tree
Showing 3 changed files with 41 additions and 32 deletions.
7 changes: 5 additions & 2 deletions src/__tests__/dataloader.test.js
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down Expand Up @@ -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();
}

Expand Down
16 changes: 8 additions & 8 deletions src/index.d.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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 {
Expand All @@ -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,
Expand Down Expand Up @@ -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>>;
}
}

Expand Down
50 changes: 28 additions & 22 deletions src/index.js
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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).
Expand All @@ -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(
Expand All @@ -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>;

/**
Expand All @@ -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) {
Expand Down Expand Up @@ -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;
}
Expand All @@ -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;
}
Expand All @@ -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) {
Expand Down Expand Up @@ -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 = [];
Expand All @@ -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
Expand Down Expand Up @@ -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
) {
Expand All @@ -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();
Expand Down

0 comments on commit bf23eb9

Please sign in to comment.