Include Model on createCachingMethods loader

src/__tests__/datasource.test.js

@@ -75,17 +75,19 @@ describe('Mongoose', () => {
     expect(getCollection(UserModel).collectionName).toBe('users')
   })
 
-  test('data source', async () => {
+  test('Data Source with Model', async () => {
     const users = new Users(UserModel)
     users.initialize()
     const user = await users.findOneById(alice._id)
     expect(user.name).toBe('Alice')
+    expect(user.id).toBe(alice._id.toString())
   })
 
-  test('collection', async () => {
+  test('Data Source with Collection', async () => {
     const users = new Users(userCollection)
     users.initialize()
     const user = await users.findOneById(alice._id)
     expect(user.name).toBe('Alice')
+    expect(user.id).toBeUndefined()
   })
 })

src/cache.js

@@ -1,6 +1,6 @@
 import DataLoader from 'dataloader'
 
-import { getCollection } from './helpers'
+import { getCollection, isModel } from './helpers'
 
 // https://github.com/graphql/dataloader#batch-function
 const orderDocs = ids => docs => {
@@ -12,12 +12,23 @@ const orderDocs = ids => docs => {
 }
 
 export const createCachingMethods = ({ collection, cache }) => {
-  const loader = new DataLoader(ids =>
-    collection
-      .find({ _id: { $in: ids } })
-      .toArray()
-      .then(orderDocs(ids))
-  )
+  const loader = new DataLoader((ids) => {
+    const res = collection.find({
+      _id: {
+        $in: ids,
+      },
+    });
+
+    let promise;
+
+    if (isModel(collection)) {
+      promise = res.exec();
+    } else {
+      promise = res.toArray();
+    }
+
+    return promise.then(orderDocs(ids));
+  });
 
   const cachePrefix = `mongo-${getCollection(collection).collectionName}-`
 

src/datasource.js

@@ -28,7 +28,7 @@ class MongoDataSource extends DataSource {
     this.context = context
 
     const methods = createCachingMethods({
-      collection: this.collection,
+      collection: this.model || this.collection,
       cache: cache || new InMemoryLRUCache()
     })