Fix creation of caching methods

Author: lorensr

README.md

@@ -6,7 +6,7 @@ Apollo [data source](https://www.apollographql.com/docs/apollo-server/features/d
 npm i apollo-datasource-mongodb
 ```
 
-This package uses [DataLoader](https://github.com/graphql/dataloader) for batching and per-request memoization caching. It also optionally (if you provide a `ttl`), does shared application-level caching (using either the default Apollo `InMemoryLRUCache` or the [cache you provide to ApolloServer()](https://www.apollographql.com/docs/apollo-server/features/data-sources#using-memcachedredis-as-a-cache-storage-backend)). It does this only for these two methods, which are added to your collections:
+This package uses [DataLoader](https://github.com/graphql/dataloader) for batching and per-request memoization caching. It also optionally (if you provide a `ttl`), does shared application-level caching (using either the default Apollo `InMemoryLRUCache` or the [cache you provide to ApolloServer()](https://www.apollographql.com/docs/apollo-server/features/data-sources#using-memcachedredis-as-a-cache-storage-backend)). It does this only for these two methods:
 
 - [`findOneById(id, options)`](#findonebyid)
 - [`findManyByIds(ids, options)`](#findmanybyids)
@@ -32,19 +32,19 @@ This package uses [DataLoader](https://github.com/graphql/dataloader) for batchi
 
 ### Basic
 
-The basic setup is subclassing `MongoDataSource`, setting your collections in the constructor, and then using the [API methods](#API) on your collections:
+The basic setup is subclassing `MongoDataSource`, setting your collections in the constructor, and then using the [API methods](#API):
 
 ```js
 import { MongoDataSource } from 'apollo-datasource-mongodb'
 
 class MyMongo extends MongoDataSource {
   constructor() {
     super()
-    this.collections = [users, posts]
+    this.collections = { users, posts }
   }
 
   getUser(userId) {
-    return users.findOneById(userId)
+    return this.users.findOneById(userId)
   }
 }
 ```
@@ -58,7 +58,7 @@ class MyMongo extends MongoDataSource {
   async getPrivateUserData(userId) {
     const isAuthorized = this.context.currentUserId === userId
     if (isAuthorized) {
-      const user = await users.findOneById(userId)
+      const user = await this.users.findOneById(userId)
       return user && user.privateData
     }
   }
@@ -71,7 +71,7 @@ If you want to implement an initialize method, it must call the parent method:
 class MyMongo extends MongoDataSource {
   constructor() {
     super()
-    this.collections = [users, posts]
+    this.collections = { users, posts }
   }
 
   initialize(config) {
@@ -103,15 +103,15 @@ client.connect(e => {
 class MyMongo extends MongoDataSource {
   constructor() {
     super()
-    this.collections = [users, posts]
+    this.collections = { users, posts }
   }
 
   getUser(userId) {
-    return users.findOneById(userId)
+    return this.users.findOneById(userId)
   }
 
   getPosts(postIds) {
-    return posts.findManyByIds(postIds)
+    return this.posts.findManyByIds(postIds)
   }
 }
 
@@ -133,28 +133,28 @@ const server = new ApolloServer({
 })
 ```
 
-You might prefer to structure it as one data source per collection, in which case you'd do:
+You might prefer to structure it as one data source per collection, in which case you'd set `this.collection` instead of `this.collections`, and the [API methods](#api) would be available directly on `this`:
 
 ```js
 class Users extends MongoDataSource {
   constructor() {
     super()
-    this.collections = [users]
+    this.collection = users
   }
 
   getUser(userId) {
-    return users.findOneById(userId)
+    return this.findOneById(userId)
   }
 }
 
 class Posts extends MongoDataSource {
   constructor() {
     super()
-    this.collections = [posts]
+    this.collection = posts
   }
 
   getPosts(postIds) {
-    return posts.findManyByIds(postIds)
+    return this.findManyByIds(postIds)
   }
 }
 
@@ -189,15 +189,15 @@ const MINUTE = 60
 class MyMongo extends MongoDataSource {
   constructor() {
     super()
-    this.collections = [users, posts]
+    this.collections = { users, posts }
   }
 
   getUser(userId) {
-    return users.findOneById(userId, { ttl: MINUTE })
+    return this.users.findOneById(userId, { ttl: MINUTE })
   }
 
   updateUserName(userId, newName) {
-    users.deleteFromCacheById(userId)
+    this.users.deleteFromCacheById(userId)
     return users.updateOne({ 
       _id: userId 
     }, {
@@ -223,18 +223,18 @@ Here we also call [`deleteFromCacheById()`](#deletefromcachebyid) to remove the
 
 ### findOneById
 
-`collection.findOneById(id, { ttl })`
+`findOneById(id, { ttl })`
 
 Resolves to the found document. Uses DataLoader to load `id`. DataLoader uses `collection.find({ _id: { $in: ids } })`. Optionally caches the document if `ttl` is set (in whole seconds).
 
 ### findManyByIds
 
-`collection.findManyByIds(ids, { ttl })`
+`findManyByIds(ids, { ttl })`
 
 Calls [`findOneById()`](#findonebyid) for each id. Resolves to an array of documents.
 
 ### deleteFromCacheById
 
-`collection.deleteFromCacheById(id)`
+`deleteFromCacheById(id)`
 
 Deletes a document from the cache.

src/__tests__/cache.test.js

@@ -1,7 +1,7 @@
 import { InMemoryLRUCache } from 'apollo-server-caching'
 import wait from 'waait'
 
-import { setupCaching } from '../cache'
+import { createCachingMethods } from '../cache'
 
 const docs = {
   id1: {
@@ -18,8 +18,9 @@ const docs = {
 const collectionName = 'test'
 const cacheKey = id => `mongo-${collectionName}-${id}`
 
-describe('setupCaching', () => {
+describe('createCachingMethods', () => {
   let collection
+  let api
   let cache
 
   beforeEach(() => {
@@ -35,23 +36,23 @@ describe('setupCaching', () => {
 
     cache = new InMemoryLRUCache()
 
-    setupCaching({ collection, cache })
+    api = createCachingMethods({ collection, cache })
   })
 
   it('adds the right methods', () => {
-    expect(collection.findOneById).toBeDefined()
-    expect(collection.findManyByIds).toBeDefined()
-    expect(collection.deleteFromCacheById).toBeDefined()
+    expect(api.findOneById).toBeDefined()
+    expect(api.findManyByIds).toBeDefined()
+    expect(api.deleteFromCacheById).toBeDefined()
   })
 
   it('finds one', async () => {
-    const doc = await collection.findOneById('id1')
+    const doc = await api.findOneById('id1')
     expect(doc).toBe(docs.id1)
     expect(collection.find.mock.calls.length).toBe(1)
   })
 
   it('finds two with batching', async () => {
-    const foundDocs = await collection.findManyByIds(['id2', 'id3'])
+    const foundDocs = await api.findManyByIds(['id2', 'id3'])
 
     expect(foundDocs[0]).toBe(docs.id2)
     expect(foundDocs[1]).toBe(docs.id3)
@@ -61,43 +62,43 @@ describe('setupCaching', () => {
 
   // TODO why doesn't this pass?
   // it.only(`doesn't cache without ttl`, async () => {
-  //   await collection.findOneById('id1')
-  //   await collection.findOneById('id1')
+  //   await api.findOneById('id1')
+  //   await api.findOneById('id1')
 
   //   expect(collection.find.mock.calls.length).toBe(2)
   // })
 
   it(`doesn't cache without ttl`, async () => {
-    await collection.findOneById('id1')
+    await api.findOneById('id1')
 
     const value = await cache.get(cacheKey('id1'))
     expect(value).toBeUndefined()
   })
 
   it(`caches`, async () => {
-    await collection.findOneById('id1', { ttl: 1 })
+    await api.findOneById('id1', { ttl: 1 })
     const value = await cache.get(cacheKey('id1'))
     expect(value).toBe(docs.id1)
 
-    await collection.findOneById('id1')
+    await api.findOneById('id1')
     expect(collection.find.mock.calls.length).toBe(1)
   })
 
   it(`caches with ttl`, async () => {
-    await collection.findOneById('id1', { ttl: 1 })
+    await api.findOneById('id1', { ttl: 1 })
     await wait(1001)
 
     const value = await cache.get(cacheKey('id1'))
     expect(value).toBeUndefined()
   })
 
   it(`deletes from cache`, async () => {
-    await collection.findOneById('id1', { ttl: 1 })
+    await api.findOneById('id1', { ttl: 1 })
 
     const valueBefore = await cache.get(cacheKey('id1'))
     expect(valueBefore).toBe(docs.id1)
 
-    await collection.deleteFromCacheById('id1')
+    await api.deleteFromCacheById('id1')
 
     const valueAfter = await cache.get(cacheKey('id1'))
     expect(valueAfter).toBeUndefined()

src/__tests__/datasource.test.js

@@ -6,19 +6,32 @@ const posts = {}
 class MyMongo extends MongoDataSource {
   constructor() {
     super()
-    this.collections = [users, posts]
+    this.collections = { users, posts }
   }
 
   initialize(config) {
     super.initialize(config)
   }
 }
 
+class SingleCollection extends MongoDataSource {
+  constructor() {
+    super()
+    this.collection = users
+  }
+}
+
 describe('MongoDataSource', () => {
   it('sets up caching functions', () => {
     const source = new MyMongo()
     source.initialize({})
-    expect(users.findOneById).toBeDefined()
-    expect(posts.findOneById).toBeDefined()
+    expect(source.users.findOneById).toBeDefined()
+    expect(source.posts.findOneById).toBeDefined()
+  })
+
+  it('sets up caching functions for single collection', () => {
+    const source = new SingleCollection()
+    source.initialize({})
+    expect(source.findOneById).toBeDefined()
   })
 })

src/cache.js

@@ -1,6 +1,6 @@
 import DataLoader from 'dataloader'
 
-export const setupCaching = ({ collection, cache }) => {
+export const createCachingMethods = ({ collection, cache }) => {
   const loader = new DataLoader(ids =>
     collection
       .find({ _id: { $in: ids } })
@@ -16,26 +16,28 @@ export const setupCaching = ({ collection, cache }) => {
 
   const cachePrefix = `mongo-${collection.collectionName}-`
 
-  collection.findOneById = async (id, { ttl } = {}) => {
-    const key = cachePrefix + id
-
-    const cacheDoc = await cache.get(key)
-    if (cacheDoc) {
-      return cacheDoc
-    }
-
-    const doc = await loader.load(id)
-    if (Number.isInteger(ttl)) {
-      // https://github.com/apollographql/apollo-server/tree/master/packages/apollo-server-caching#apollo-server-caching
-      cache.set(key, doc, { ttl })
-    }
-
-    return doc
-  }
-
-  collection.findManyByIds = (ids, { ttl } = {}) => {
-    return Promise.all(ids.map(id => collection.findOneById(id, { ttl })))
+  const methods = {
+    findOneById: async (id, { ttl } = {}) => {
+      const key = cachePrefix + id
+
+      const cacheDoc = await cache.get(key)
+      if (cacheDoc) {
+        return cacheDoc
+      }
+
+      const doc = await loader.load(id)
+      if (Number.isInteger(ttl)) {
+        // https://github.com/apollographql/apollo-server/tree/master/packages/apollo-server-caching#apollo-server-caching
+        cache.set(key, doc, { ttl })
+      }
+
+      return doc
+    },
+    findManyByIds: (ids, { ttl } = {}) => {
+      return Promise.all(ids.map(id => methods.findOneById(id, { ttl })))
+    },
+    deleteFromCacheById: id => cache.delete(cachePrefix + id)
   }
 
-  collection.deleteFromCacheById = id => cache.delete(cachePrefix + id)
+  return methods
 }

src/datasource.js

@@ -2,22 +2,37 @@ import { DataSource } from 'apollo-datasource'
 import { ApolloError } from 'apollo-server-errors'
 import { InMemoryLRUCache } from 'apollo-server-caching'
 
-import { setupCaching } from './cache'
+import { createCachingMethods } from './cache'
 
 class MongoDataSource extends DataSource {
   // https://github.com/apollographql/apollo-server/blob/master/packages/apollo-datasource/src/index.ts
   initialize(config) {
     this.context = config.context
 
-    if (!this.collections || !this.collections.length) {
+    const setUpCorrectly =
+      typeof this.collections === 'object' || this.collection
+    if (!setUpCorrectly) {
       throw new ApolloError(
-        'Child class of MongoDataSource must set this.collections in constructor'
+        'Child class of MongoDataSource must set this.collections or this.collection in constructor'
       )
     }
 
     const cache = config.cache || new InMemoryLRUCache()
 
-    this.collections.forEach(collection => setupCaching({ collection, cache }))
+    if (this.collections) {
+      for (const key in this.collections) {
+        this[key] = createCachingMethods({
+          collection: this.collections[key],
+          cache
+        })
+      }
+    } else {
+      const methods = createCachingMethods({
+        collection: this.collection,
+        cache
+      })
+      Object.assign(this, methods)
+    }
   }
 }