fix(cursor): set readPreference for cursor.count

Fixes NODE-1581
mongodb · Aug 1, 2018 · 13d776f · 13d776f
1 parent a5d0f1d
commit 13d776f
Show file tree

Hide file tree

Showing 3 changed files with 103 additions and 51 deletions.
diff --git a/lib/operations/collection_ops.js b/lib/operations/collection_ops.js
@@ -175,33 +175,14 @@ function checkForAtomicOperators(update) {
 function count(coll, query, options, callback) {
   if (typeof options === 'function') (callback = options), (options = {});
   options = Object.assign({}, options);
+  options.collectionName = coll.s.name;
 
-  const skip = options.skip;
-  const limit = options.limit;
-  const hint = options.hint;
-  const maxTimeMS = options.maxTimeMS;
-  query = query || {};
-
-  // Final query
-  const cmd = {
-    count: coll.s.name,
-    query: query
-  };
-
-  // Add limit, skip and maxTimeMS if defined
-  if (typeof skip === 'number') cmd.skip = skip;
-  if (typeof limit === 'number') cmd.limit = limit;
-  if (typeof maxTimeMS === 'number') cmd.maxTimeMS = maxTimeMS;
-  if (hint) cmd.hint = hint;
-
-  // Ensure we have the right read preference inheritance
-  options.readPreference = resolveReadPreference(options, { db: coll.s.db, collection: coll });
-
-  // Do we have a readConcern specified
-  decorateWithReadConcern(cmd, coll, options);
+  options.readPreference = resolveReadPreference(options, {
+    db: coll.s.db,
+    collection: coll
+  });
 
-  // Have we specified collation
-  decorateWithCollation(cmd, coll, options);
+  const cmd = buildCountCommand(coll, query, options);
 
   executeCommand(coll.s.db, cmd, options, (err, result) => {
     if (err) return handleCallback(callback, err);
@@ -236,6 +217,51 @@ function countDocuments(coll, query, options, callback) {
   });
 }
 
+/**
+ * Build the count command.
+ *
+ * @method
+ * @param {collectionOrCursor} an instance of a collection or cursor
+ * @param {object} query The query for the count.
+ * @param {object} [options] Optional settings. See Collection.prototype.count and Cursor.prototype.count for a list of options.
+ */
+function buildCountCommand(collectionOrCursor, query, options) {
+  const skip = options.skip;
+  const limit = options.limit;
+  let hint = options.hint;
+  const maxTimeMS = options.maxTimeMS;
+  query = query || {};
+
+  // Final query
+  const cmd = {
+    count: options.collectionName,
+    query: query
+  };
+
+  // check if collectionOrCursor is a cursor by using cursor.s.numberOfRetries
+  if (collectionOrCursor.s.numberOfRetries) {
+    if (collectionOrCursor.s.options.hint) {
+      hint = collectionOrCursor.s.options.hint;
+    } else if (collectionOrCursor.s.cmd.hint) {
+      hint = collectionOrCursor.s.cmd.hint;
+    }
+    decorateWithCollation(cmd, collectionOrCursor, collectionOrCursor.s.cmd);
+  } else {
+    decorateWithCollation(cmd, collectionOrCursor, options);
+  }
+
+  // Add limit, skip and maxTimeMS if defined
+  if (typeof skip === 'number') cmd.skip = skip;
+  if (typeof limit === 'number') cmd.limit = limit;
+  if (typeof maxTimeMS === 'number') cmd.maxTimeMS = maxTimeMS;
+  if (hint) cmd.hint = hint;
+
+  // Do we have a readConcern specified
+  decorateWithReadConcern(cmd, collectionOrCursor);
+
+  return cmd;
+}
+
 /**
  * Create an index on the db and collection.
  *
@@ -1391,6 +1417,7 @@ module.exports = {
   checkForAtomicOperators,
   count,
   countDocuments,
+  buildCountCommand,
   createIndex,
   createIndexes,
   deleteMany,

diff --git a/lib/operations/cursor_ops.js b/lib/operations/cursor_ops.js
@@ -1,6 +1,6 @@
 'use strict';
 
-const decorateWithCollation = require('../utils').decorateWithCollation;
+const buildCountCommand = require('./collection_ops').buildCountCommand;
 const formattedOrderClause = require('../utils').formattedOrderClause;
 const handleCallback = require('../utils').handleCallback;
 const MongoError = require('mongodb-core').MongoError;
@@ -21,37 +21,30 @@ function count(cursor, applySkipLimit, opts, callback) {
     if (typeof cursor.cursorLimit() === 'number') opts.limit = cursor.cursorLimit();
   }
 
-  // Command
-  const delimiter = cursor.s.ns.indexOf('.');
-
-  let command = {
-    count: cursor.s.ns.substr(delimiter + 1),
-    query: cursor.s.cmd.query
-  };
-
-  // Apply a readConcern if set
-  if (cursor.s.cmd.readConcern) {
-    command.readConcern = cursor.s.cmd.readConcern;
+  // Ensure we have the right read preference inheritance
+  if (opts.readPreference) {
+    cursor.setReadPreference(opts.readPreference);
   }
 
-  // Apply a hint if set
-  if (cursor.s.cmd.hint) {
-    command.hint = cursor.s.cmd.hint;
+  if (
+    typeof opts.maxTimeMS !== 'number' &&
+    cursor.s.cmd &&
+    typeof cursor.s.cmd.maxTimeMS === 'number'
+  ) {
+    opts.maxTimeMS = cursor.s.cmd.maxTimeMS;
   }
 
-  // Apply a collation if set
-  decorateWithCollation(command, cursor, cursor.s.cmd);
+  let options = {};
+  options.skip = opts.skip;
+  options.limit = opts.limit;
+  options.hint = opts.hint;
+  options.maxTimeMS = opts.maxTimeMS;
 
-  if (typeof opts.maxTimeMS === 'number') {
-    command.maxTimeMS = opts.maxTimeMS;
-  } else if (cursor.s.cmd && typeof cursor.s.cmd.maxTimeMS === 'number') {
-    command.maxTimeMS = cursor.s.cmd.maxTimeMS;
-  }
+  // Command
+  const delimiter = cursor.s.ns.indexOf('.');
+  options.collectionName = cursor.s.ns.substr(delimiter + 1);
 
-  // Merge in any options
-  if (opts.skip) command.skip = opts.skip;
-  if (opts.limit) command.limit = opts.limit;
-  if (cursor.s.options.hint) command.hint = cursor.s.options.hint;
+  const command = buildCountCommand(cursor, cursor.s.cmd.query, options);
 
   // Set cursor server to the same as the topology
   cursor.server = cursor.topology.s.coreTopology;

diff --git a/test/functional/cursor_tests.js b/test/functional/cursor_tests.js
@@ -5,6 +5,7 @@ const fs = require('fs');
 const expect = require('chai').expect;
 const Long = require('bson').Long;
 const sinon = require('sinon');
+const ReadPreference = require('mongodb-core').ReadPreference;
 
 describe('Cursor', function() {
   before(function() {
@@ -280,6 +281,37 @@ describe('Cursor', function() {
     }
   });
 
+  it('Should correctly execute cursor count with secondary readPreference', {
+    // Add a tag that our runner can trigger on
+    // in this case we are setting that node needs to be higher than 0.10.X to run
+    metadata: {
+      requires: { topology: 'replicaset' }
+    },
+
+    // The actual test we wish to run
+    test: function(done) {
+      const configuration = this.configuration;
+      const client = configuration.newClient(configuration.writeConcernMax(), { poolSize: 1 });
+
+      client.connect((err, client) => {
+        const db = client.db(configuration.db);
+
+        const internalClientCursor = sinon.spy(client.topology.s.coreTopology, 'cursor');
+        const expectedReadPreference = new ReadPreference(ReadPreference.SECONDARY);
+
+        const cursor = db.collection('countTEST').find({ qty: { $gt: 4 } });
+        cursor.count(true, { readPreference: ReadPreference.SECONDARY }, err => {
+          expect(err).to.be.null;
+          expect(internalClientCursor.getCall(0).args[2])
+            .to.have.nested.property('readPreference')
+            .that.deep.equals(expectedReadPreference);
+          client.close();
+          done();
+        });
+      });
+    }
+  });
+
   /**
    * @ignore
    * @api private