Reinstated obedient commands for DB

JAVA-1784

Author: rozza

driver/src/main/com/mongodb/DB.java

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2008-2014 MongoDB, Inc.
+ * Copyright 2008-2015 MongoDB, Inc.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -35,6 +35,7 @@
 
 import java.util.ArrayList;
 import java.util.Collections;
+import java.util.HashSet;
 import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Set;
@@ -355,7 +356,7 @@ public CommandResult command(final DBObject command, final DBEncoder encoder) {
      */
     public CommandResult command(final DBObject command, final ReadPreference readPreference, final DBEncoder encoder) {
         try {
-            return executeCommand(wrap(command, encoder), readPreference);
+            return executeCommand(wrap(command, encoder), getCommandReadPreference(command, readPreference));
         } catch (MongoCommandException ex) {
             return new CommandResult(ex.getResponse(), ex.getServerAddress());
         }
@@ -617,4 +618,43 @@ private BsonDocument wrap(final DBObject document, final DBEncoder encoder) {
             return new BsonDocumentWrapper<DBObject>(document, new DBEncoderAdapter(encoder));
         }
     }
+
+    /**
+     * Determines the read preference that should be used for the given command.
+     *
+     * @param command             the {@link DBObject} representing the command
+     * @param requestedPreference the preference requested by the client.
+     * @return the read preference to use for the given command.  It will never return {@code null}.
+     * @see com.mongodb.ReadPreference
+     */
+    ReadPreference getCommandReadPreference(final DBObject command, final ReadPreference requestedPreference) {
+        String comString = command.keySet().iterator().next().toLowerCase();
+        boolean primaryRequired = !OBEDIENT_COMMANDS.contains(comString);
+
+        if (primaryRequired) {
+            return ReadPreference.primary();
+        } else if (requestedPreference == null) {
+            return ReadPreference.primary();
+        } else {
+            return requestedPreference;
+        }
+    }
+
+    private static final Set<String> OBEDIENT_COMMANDS = new HashSet<String>();
+
+    static {
+        OBEDIENT_COMMANDS.add("aggregate");
+        OBEDIENT_COMMANDS.add("collstats");
+        OBEDIENT_COMMANDS.add("count");
+        OBEDIENT_COMMANDS.add("dbstats");
+        OBEDIENT_COMMANDS.add("distinct");
+        OBEDIENT_COMMANDS.add("geonear");
+        OBEDIENT_COMMANDS.add("geosearch");
+        OBEDIENT_COMMANDS.add("geowalk");
+        OBEDIENT_COMMANDS.add("group");
+        OBEDIENT_COMMANDS.add("listcollections");
+        OBEDIENT_COMMANDS.add("listindexes");
+        OBEDIENT_COMMANDS.add("parallelcollectionscan");
+        OBEDIENT_COMMANDS.add("text");
+    }
 }

driver/src/test/unit/com/mongodb/DBSpecification.groovy

@@ -1,7 +1,24 @@
+/*
+ * Copyright 2014-2015 MongoDB, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
 package com.mongodb
 
 import com.mongodb.operation.CreateCollectionOperation
 import org.bson.BsonDocument
+import org.bson.BsonDouble
 import spock.lang.Specification
 
 import static com.mongodb.CustomMatchers.isTheSameAs
@@ -35,4 +52,53 @@ class DBSpecification extends Specification {
                                               .autoIndex(true)
                                               .storageEngineOptions(new BsonDocument('wiredTiger', new BsonDocument())))
     }
-}
+
+
+    def 'should use provided read preference for obedient commands'() {
+        given:
+        def mongo = Stub(Mongo)
+        mongo.mongoClientOptions >> MongoClientOptions.builder().build()
+        def executor = new TestOperationExecutor([new BsonDocument('ok', new BsonDouble(1.0))])
+        def database = new DB(mongo, 'test', executor)
+        database.setReadPreference(ReadPreference.secondary())
+
+        when:
+        database.command(cmd)
+
+        then:
+        executor.getReadPreference() == expectedReadPreference
+
+        where:
+        expectedReadPreference     | cmd
+        ReadPreference.secondary() | new BasicDBObject('listCollections', 1)
+        ReadPreference.secondary() | new BasicDBObject('collStats', 1)
+        ReadPreference.secondary() | new BasicDBObject('dbStats', 1)
+        ReadPreference.secondary() | new BasicDBObject('distinct', 1)
+        ReadPreference.secondary() | new BasicDBObject('geoNear', 1)
+        ReadPreference.secondary() | new BasicDBObject('geoSearch', 1)
+        ReadPreference.secondary() | new BasicDBObject('group', 1)
+        ReadPreference.secondary() | new BasicDBObject('listCollections', 1)
+        ReadPreference.secondary() | new BasicDBObject('listIndexes', 1)
+        ReadPreference.secondary() | new BasicDBObject('parallelCollectionScan', 1)
+        ReadPreference.secondary() | new BasicDBObject('text', 1)
+    }
+
+    def 'should use primary read preference for non obedient commands'() {
+        given:
+        def mongo = Stub(Mongo)
+        mongo.mongoClientOptions >> MongoClientOptions.builder().build()
+        def executor = new TestOperationExecutor([new BsonDocument('ok', new BsonDouble(1.0))])
+        def database = new DB(mongo, 'test', executor)
+        database.setReadPreference(ReadPreference.secondary())
+
+        when:
+        database.command(cmd)
+
+        then:
+        executor.getReadPreference() == expectedReadPreference
+
+        where:
+        expectedReadPreference      | cmd
+        ReadPreference.primary()    | new BasicDBObject('command', 1)
+    }
+}