SERVER-7423 Support commands for read preference

Step#2 Added mongos support.

1 files changed, 112 insertions, 74 deletions

diff --git a/jstests/sharding/read_pref_cmd.js b/jstests/sharding/read_pref_cmd.js
index 763bdcbc4dd..0b271737eb8 100644
--- a/jstests/sharding/read_pref_cmd.js
+++ b/jstests/sharding/read_pref_cmd.js
@@ -1,96 +1,133 @@
-var NUM_NODES = 2;
+var NODE_COUNT = 2;
 
 /**
- * Performs a series of test on commands with read preference
+ * Performs a series of tests on commands with read preference.
  *
  * @param conn {Mongo} the connection object of which to test the read
  *     preference functionality.
  * @param hostList {Array.<Mongo>} list of the replica set host members.
+ * @param isMongos {boolean} true if conn is a mongos connection.
  */
-var doTest = function(conn, hostList) {
+var doTest = function(conn, hostList, isMongos) {
     var testDB = conn.getDB('test');
     conn.setSlaveOk(false); // purely rely on readPref
     conn.setReadPref('secondary');
 
-    // Test command that can be sent to secondary
-    var distinctResult = testDB.runCommand({ distinct: 'user',
-            key: { x: 1 }, query: { x: 1 }});
-
-    var testedAtLeastOnce = false;
-    hostList.forEach(function(node) {
-        var testDB = node.getDB('test');
-        var result = testDB.system.profile.findOne({ op: 'command',
-            ns: 'test.$cmd', 'command.query.distinct': 'user' });
-
-        if (result != null) {
-            assert(testDB.adminCommand({ isMaster: 1 }).secondary);
-            testedAtLeastOnce = true;
+    /**
+     * Performs the command and checks whether the command was routed to the
+     * appropriate node.
+     *
+     * @param cmdObj the cmd to send.
+     * @param secOk true if command should be routed to a secondary.
+     * @param profileQuery the query to perform agains the profile collection to
+     *     look for the cmd just sent.
+     */
+    var cmdTest = function(cmdObj, secOk, profileQuery) {
+        var cmdResult = testDB.runCommand(cmdObj, { readPref: { mode: 'secondary' }});
+        jsTest.log('cmd result: ' + tojson(cmdResult));
+        assert(cmdResult.ok);
+
+        var testedAtLeastOnce = false;
+        var query = { op: 'command', ns: 'test.$cmd' };
+        Object.extend(query, profileQuery);
+
+        hostList.forEach(function(node) {
+            var testDB = node.getDB('test');
+            var result = testDB.system.profile.findOne(query);
+
+            if (result != null) {
+                if (secOk) {
+                    assert(testDB.adminCommand({ isMaster: 1 }).secondary);
+                }
+                else {
+                    assert(testDB.adminCommand({ isMaster: 1 }).ismaster);
+                }
+
+                testedAtLeastOnce = true;
+            }
+        });
+
+        assert(testedAtLeastOnce);
+    };
+
+    /**
+     * Assumption: all values are native types (no objects)
+     */
+    var formatProfileQuery = function(queryObj, isEmbedded) {
+        var prefix = isEmbedded? 'command.query.' : 'command.';
+        var newObj = {};
+
+        for (var field in queryObj) {
+            newObj[prefix + field] = queryObj[field];
         }
-    });
 
-    assert(testedAtLeastOnce);
+        return newObj;
+    };
 
-    // Test command that can't be sent to secondary
-    var createResult = testDB.runCommand({ create: 'user' });
-    assert(createResult.ok, 'create cmd failed: ' + tojson(createResult));
-    testDB.runCommand({ getLastError: 1, w: NUM_NODES });
-
-    testedAtLeastOnce = false;
-    hostList.forEach(function(node) {
-        var testDB = node.getDB('test');
-        var result = testDB.system.profile.findOne({ op: 'command',
-            ns: 'test.$cmd', 'command.query.create': 'user' });
-
-        if (result != null) {
-            assert(testDB.adminCommand({ isMaster: 1 }).ismaster);
-            testedAtLeastOnce = true;
-        }
-    });
+    // Test command that can be sent to secondary
+    cmdTest({ distinct: 'user', key: { x: 1 }, query: { x: 1 }}, true,
+        formatProfileQuery({ distinct: 'user' }, !isMongos));
 
-    assert(testedAtLeastOnce);
+    // Test command that can't be sent to secondary
+    cmdTest({ create: 'mrIn' }, false, formatProfileQuery({ create: 'mrIn' }, !isMongos));
+    // Make sure mrIn is propagated to secondaries before proceeding
+    testDB.runCommand({ getLastError: 1, w: NODE_COUNT });
 
-    // Test inline map reduce
     var mapFunc = function(doc) {};
     var reduceFunc = function(key, values) { return values; };
 
-    var inlineMRResult = testDB.runCommand({ mapreduce: 'user', map: mapFunc,
-        reduce: reduceFunc, out: { inline: 1 }});
-    assert(inlineMRResult.ok, 'inline mr failed: ' + tojson(inlineMRResult));
+    // Test inline mapReduce on sharded collection.
+    // Note that in sharded map reduce, it will output the result in a temp collection
+    // even if out is inline.
+    if (isMongos) {
+        cmdTest({ mapreduce: 'user', map: mapFunc, reduce: reduceFunc, out: { inline: 1 }},
+            false, formatProfileQuery({ mapreduce: 'user', shardedFirstPass: true }, false));
+    }
 
-    testedAtLeastOnce = false;
-    hostList.forEach(function(node) {
-        var testDB = node.getDB('test');
-        var result = testDB.system.profile.findOne({ op: 'command',
-            ns: 'test.$cmd', 'command.query.mapreduce': 'user' });
+    // Test inline mapReduce on unsharded collection.
+    cmdTest({ mapreduce: 'mrIn', map: mapFunc, reduce: reduceFunc, out: { inline: 1 }}, true,
+        formatProfileQuery({ mapreduce: 'mrIn', 'out.inline': 1 }, !isMongos));
 
-        if (result != null) {
-            assert(testDB.adminCommand({ isMaster: 1 }).secondary);
-            testedAtLeastOnce = true;
-        }
-    });
-
-    // Test non-inline map reduce
-    testDB.runCommand({ create: 'mrIn' });
-    testDB.runCommand({ getLastError: 1, w: NUM_NODES });
-    var outCollMRResult = testDB.runCommand({ mapreduce: 'mrIn', map: mapFunc,
-        reduce: reduceFunc, out: { replace: 'mrOut' }});
-    assert(outCollMRResult.ok, 'replace mr cmd failed: ' + tojson(outCollMRResult));
-
-    testedAtLeastOnce = false;
-    hostList.forEach(function(node) {
-        var testDB = node.getDB('test');
-        var result = testDB.system.profile.findOne({ op: 'command',
-            ns: 'test.$cmd', 'command.query.mapreduce': 'mrIn' });
-
-        if (result != null) {
-            assert(testDB.adminCommand({ isMaster: 1 }).ismaster);
-            testedAtLeastOnce = true;
-        }
-    });
+    // Test non-inline mapReduce on sharded collection.
+    if (isMongos) {
+        cmdTest({ mapreduce: 'user', map: mapFunc, reduce: reduceFunc,
+            out: { replace: 'mrOut' }}, false,
+            formatProfileQuery({ mapreduce: 'user', shardedFirstPass: true }, false));
+    }
+
+    // Test non-inline mapReduce on unsharded collection.
+    cmdTest({ mapreduce: 'mrIn', map: mapFunc, reduce: reduceFunc, out: { replace: 'mrOut' }},
+        false, formatProfileQuery({ mapreduce: 'mrIn', 'out.replace': 'mrOut' }, !isMongos));
+
+    // Test other commands that can be sent to secondary.
+    cmdTest({ count: 'user' }, true, formatProfileQuery({ count: 'user' }, !isMongos));
+    cmdTest({ group: { key: { x: true }, '$reduce': function(a, b) {}, ns: 'mrIn',
+        initial: { x: 0  }}}, true, formatProfileQuery({ 'group.ns': 'mrIn' }, !isMongos));
+
+    cmdTest({ collStats: 'user' }, true, formatProfileQuery({ count: 'user' }, !isMongos));
+    cmdTest({ dbStats: 1 }, true, formatProfileQuery({ dbStats: 1 }, !isMongos));
+
+    testDB.user.ensureIndex({ loc: '2d' });
+    testDB.runCommand({ getLastError: 1, w: NODE_COUNT });
+    cmdTest({ geoNear: 'user', near: [1, 1] }, true,
+        formatProfileQuery({ geoNear: 'user' }, !isMongos));
+
+    // Test on sharded
+    cmdTest({ aggregate: 'user', pipeline: [{ $project: { x: 1 }}] }, true,
+        formatProfileQuery({ aggregate: 'user' }, !isMongos));
+
+    // Test on non-sharded
+    cmdTest({ aggregate: 'mrIn', pipeline: [{ $project: { x: 1 }}] }, true,
+        formatProfileQuery({ aggregate: 'mrIn' }, !isMongos));
 };
 
-var st = new ShardingTest({ shards: { rs0: { nodes: NUM_NODES }}});
+var st = new ShardingTest({ shards: { rs0: { nodes: NODE_COUNT }}});
 st.stopBalancer();
+
+var configDB = st.s.getDB('config');
+configDB.adminCommand({ enableSharding: 'test' });
+configDB.adminCommand({ shardCollection: 'test.user', key: { x: 1 }});
+
 ReplSetTest.awaitRSClientHosts(st.s, st.rs0.nodes);
 
 st.rs0.nodes.forEach(function(node) {
@@ -115,19 +152,20 @@ assert.soon(function() {
     }
 });
 
-doTest(replConn, st.rs0.nodes);
+doTest(new Mongo(st.rs0.getURL()), st.rs0.nodes, false);
 
-// TODO: uncomment once read preference command is properly implemented in mongos
-/*
 st.s.getDB('test').dropDatabase();
+
+configDB.adminCommand({ enableSharding: 'test' });
+configDB.adminCommand({ shardCollection: 'test.user', key: { x: 1 }});
+
 st.rs0.nodes.forEach(function(node) {
     node.getDB('test').setProfilingLevel(2);
 });
 
 jsTest.log('Starting test for mongos connection');
 
-doTest(st.s, st.rs0.nodes);
-*/
+doTest(st.s, st.rs0.nodes, true);
 
 st.stop();