SERVER-29319 Expose OpMsgRequest interface in DBClient

11 files changed, 117 insertions, 190 deletions

diff --git a/src/mongo/client/dbclient.cpp b/src/mongo/client/dbclient.cpp
index 669ca102a6a..4732658795c 100644
--- a/src/mongo/client/dbclient.cpp
+++ b/src/mongo/client/dbclient.cpp
@@ -169,13 +169,12 @@ const rpc::ReplyMetadataReader& DBClientWithCommands::getReplyMetadataReader() {
     return _metadataReader;
 }
 
-rpc::UniqueReply DBClientWithCommands::runCommandWithMetadata(StringData database,
-                                                              StringData command,
-                                                              const BSONObj& metadataIn,
-                                                              const BSONObj& commandArgs) {
+std::pair<rpc::UniqueReply, DBClientWithCommands*> DBClientWithCommands::runCommandWithTarget(
+    OpMsgRequest request) {
     uassert(ErrorCodes::InvalidNamespace,
-            str::stream() << "Database name '" << database << "' is not valid.",
-            NamespaceString::validDBName(database, NamespaceString::DollarInDbNameBehavior::Allow));
+            str::stream() << "Database name '" << request.getDatabase() << "' is not valid.",
+            NamespaceString::validDBName(request.getDatabase(),
+                                         NamespaceString::DollarInDbNameBehavior::Allow));
 
     // Make sure to reconnect if needed before building our request, since the request depends on
     // the negotiated protocol which can change due to a reconnect.
@@ -184,19 +183,15 @@ rpc::UniqueReply DBClientWithCommands::runCommandWithMetadata(StringData databas
     // call() oddly takes this by pointer, so we need to put it on the stack.
     auto host = getServerAddress();
 
-    auto metadata = metadataIn;
-
     if (_metadataWriter) {
-        BSONObjBuilder metadataBob(std::move(metadata));
+        BSONObjBuilder metadataBob(std::move(request.body));
         uassertStatusOK(
             _metadataWriter((haveClient() ? cc().getOperationContext() : nullptr), &metadataBob));
-        metadata = metadataBob.obj();
+        request.body = metadataBob.obj();
     }
 
-    auto requestMsg = rpc::messageFromOpMsgRequest(
-        getClientRPCProtocols(),
-        getServerRPCProtocols(),
-        OpMsgRequest::fromDBAndBody(database, std::move(commandArgs), metadata));
+    auto requestMsg =
+        rpc::messageFromOpMsgRequest(getClientRPCProtocols(), getServerRPCProtocols(), request);
 
     Message replyMsg;
 
@@ -206,7 +201,7 @@ rpc::UniqueReply DBClientWithCommands::runCommandWithMetadata(StringData databas
     uassert(ErrorCodes::HostUnreachable,
             str::stream() << "network error while attempting to run "
                           << "command '"
-                          << command
+                          << request.getCommandName()
                           << "' "
                           << "on host '"
                           << host
@@ -234,43 +229,48 @@ rpc::UniqueReply DBClientWithCommands::runCommandWithMetadata(StringData databas
                                        commandReply->getCommandReply());
     }
 
-    return rpc::UniqueReply(std::move(replyMsg), std::move(commandReply));
+    return {rpc::UniqueReply(std::move(replyMsg), std::move(commandReply)), this};
+}
+
+rpc::UniqueReply DBClientWithCommands::runCommandWithMetadata(StringData database,
+                                                              StringData command,
+                                                              const BSONObj& metadata,
+                                                              BSONObj commandArgs) {
+    return runCommand(OpMsgRequest::fromDBAndBody(database, std::move(commandArgs), metadata));
 }
 
 std::tuple<rpc::UniqueReply, DBClientWithCommands*>
 DBClientWithCommands::runCommandWithMetadataAndTarget(StringData database,
                                                       StringData command,
                                                       const BSONObj& metadata,
-                                                      const BSONObj& commandArgs) {
-    return std::make_tuple(runCommandWithMetadata(database, command, metadata, commandArgs), this);
+                                                      BSONObj commandArgs) {
+    return runCommandWithTarget(
+        OpMsgRequest::fromDBAndBody(database, std::move(commandArgs), metadata));
 }
 
 std::tuple<bool, DBClientWithCommands*> DBClientWithCommands::runCommandWithTarget(
-    const string& dbname, const BSONObj& cmd, BSONObj& info, int options) {
+    const string& dbname, BSONObj cmd, BSONObj& info, int options) {
     BSONObj upconvertedCmd;
     BSONObj upconvertedMetadata;
 
     // TODO: This will be downconverted immediately if the underlying
     // requestBuilder is a legacyRequest builder. Not sure what the best
     // way to get around that is without breaking the abstraction.
-    std::tie(upconvertedCmd, upconvertedMetadata) = rpc::upconvertRequestMetadata(cmd, options);
-
-    auto commandName = upconvertedCmd.firstElementFieldName();
+    std::tie(upconvertedCmd, upconvertedMetadata) =
+        rpc::upconvertRequestMetadata(std::move(cmd), options);
 
-    auto resultTuple =
-        runCommandWithMetadataAndTarget(dbname, commandName, upconvertedMetadata, upconvertedCmd);
-    auto result = std::move(std::get<0>(resultTuple));
+    auto result = runCommandWithTarget(
+        OpMsgRequest::fromDBAndBody(dbname, std::move(upconvertedCmd), upconvertedMetadata));
 
-    info = result->getCommandReply().getOwned();
-
-    return std::make_tuple(isOk(info), std::get<1>(resultTuple));
+    info = result.first->getCommandReply().getOwned();
+    return std::make_tuple(isOk(info), result.second);
 }
 
 bool DBClientWithCommands::runCommand(const string& dbname,
-                                      const BSONObj& cmd,
+                                      BSONObj cmd,
                                       BSONObj& info,
                                       int options) {
-    auto res = runCommandWithTarget(dbname, cmd, info, options);
+    auto res = runCommandWithTarget(dbname, std::move(cmd), info, options);
     return std::get<0>(res);
 }
 
@@ -918,18 +918,18 @@ void DBClientConnection::logout(const string& dbname, BSONObj& info) {
     runCommand(dbname, BSON("logout" << 1), info);
 }
 
-bool DBClientConnection::runCommand(const string& dbname,
-                                    const BSONObj& cmd,
-                                    BSONObj& info,
-                                    int options) {
-    if (DBClientWithCommands::runCommand(dbname, cmd, info, options))
-        return true;
+std::pair<rpc::UniqueReply, DBClientWithCommands*> DBClientConnection::runCommandWithTarget(
+    OpMsgRequest request) {
 
+    auto out = DBClientWithCommands::runCommandWithTarget(std::move(request));
     if (!_parentReplSetName.empty()) {
-        handleNotMasterResponse(info["errmsg"]);
+        const auto replyBody = out.first->getCommandReply();
+        if (!isOk(replyBody)) {
+            handleNotMasterResponse(replyBody["errmsg"]);
+        }
     }
 
-    return false;
+    return out;
 }
 
 void DBClientConnection::_checkConnection() {
diff --git a/src/mongo/client/dbclient_rs.cpp b/src/mongo/client/dbclient_rs.cpp
index cd80a61f79d..38948381508 100644
--- a/src/mongo/client/dbclient_rs.cpp
+++ b/src/mongo/client/dbclient_rs.cpp
@@ -903,19 +903,8 @@ void DBClientReplicaSet::checkResponse(const char* data,
     }
 }
 
-rpc::UniqueReply DBClientReplicaSet::runCommandWithMetadata(StringData database,
-                                                            StringData command,
-                                                            const BSONObj& metadata,
-                                                            const BSONObj& commandArgs) {
-    auto ret = runCommandWithMetadataAndTarget(database, command, metadata, commandArgs);
-    return std::move(std::get<0>(ret));
-}
-
-std::tuple<rpc::UniqueReply, DBClientWithCommands*>
-DBClientReplicaSet::runCommandWithMetadataAndTarget(StringData database,
-                                                    StringData command,
-                                                    const BSONObj& metadata,
-                                                    const BSONObj& commandArgs) {
+std::pair<rpc::UniqueReply, DBClientWithCommands*> DBClientReplicaSet::runCommandWithTarget(
+    OpMsgRequest request) {
     // This overload exists so we can parse out the read preference and then use server
     // selection directly without having to re-parse the raw message.
 
@@ -923,15 +912,14 @@ DBClientReplicaSet::runCommandWithMetadataAndTarget(StringData database,
     // so we don't have to re-parse it, however, that will come with its own set of
     // complications (e.g. some kind of base class or concept for MetadataSerializable
     // objects). For now we do it the stupid way.
-    auto readPref = uassertStatusOK(ReadPreferenceSetting::fromContainingBSON(metadata));
+    auto readPref = uassertStatusOK(ReadPreferenceSetting::fromContainingBSON(request.body));
 
     if (readPref.pref == ReadPreference::PrimaryOnly ||
         // If the command is not runnable on a secondary, we run it on the primary
         // regardless of the read preference.
-        !_isSecondaryCommand(command, commandArgs)) {
+        !_isSecondaryCommand(request.getCommandName(), request.body)) {
         auto conn = checkMaster();
-        return std::make_tuple(
-            conn->runCommandWithMetadata(database, command, metadata, commandArgs), conn);
+        return conn->runCommandWithTarget(std::move(request));
     }
 
     auto rpShared = std::make_shared<ReadPreferenceSetting>(std::move(readPref));
@@ -942,10 +930,8 @@ DBClientReplicaSet::runCommandWithMetadataAndTarget(StringData database,
             if (conn == nullptr) {
                 break;
             }
-            // We can't move database and command in case this throws
-            // and we retry.
-            return std::make_tuple(
-                conn->runCommandWithMetadata(database, command, metadata, commandArgs), conn);
+            // We can't move the request since we need it to retry.
+            return conn->runCommandWithTarget(request);
         } catch (const DBException& ex) {
             _invalidateLastSlaveOkCache(ex.toStatus());
         }
@@ -954,7 +940,7 @@ DBClientReplicaSet::runCommandWithMetadataAndTarget(StringData database,
     uasserted(ErrorCodes::NodeNotFound,
               str::stream() << "Could not satisfy $readPreference of '" << readPref.toString()
                             << "' while attempting to run command "
-                            << command);
+                            << request.getCommandName());
 }
 
 bool DBClientReplicaSet::call(Message& toSend,
diff --git a/src/mongo/client/dbclient_rs.h b/src/mongo/client/dbclient_rs.h
index 02e9db27f65..5369bee9ea6 100644
--- a/src/mongo/client/dbclient_rs.h
+++ b/src/mongo/client/dbclient_rs.h
@@ -186,16 +186,9 @@ public:
         return true;
     }
 
-    rpc::UniqueReply runCommandWithMetadata(StringData database,
-                                            StringData command,
-                                            const BSONObj& metadata,
-                                            const BSONObj& commandArgs) final;
-
-    std::tuple<rpc::UniqueReply, DBClientWithCommands*> runCommandWithMetadataAndTarget(
-        StringData database,
-        StringData command,
-        const BSONObj& metadata,
-        const BSONObj& commandArgs) final;
+    using DBClientWithCommands::runCommandWithTarget;
+    std::pair<rpc::UniqueReply, DBClientWithCommands*> runCommandWithTarget(
+        OpMsgRequest request) final;
 
     void setRequestMetadataWriter(rpc::RequestMetadataWriter writer) final;
 
diff --git a/src/mongo/client/dbclientinterface.h b/src/mongo/client/dbclientinterface.h
index 6d2976d08e4..41296ff6a72 100644
--- a/src/mongo/client/dbclientinterface.h
+++ b/src/mongo/client/dbclientinterface.h
@@ -47,6 +47,7 @@
 #include "mongo/util/mongoutils/str.h"
 #include "mongo/util/net/abstract_message_port.h"
 #include "mongo/util/net/message.h"
+#include "mongo/util/net/op_msg.h"
 
 namespace mongo {
 
@@ -410,15 +411,29 @@ public:
     const rpc::ReplyMetadataReader& getReplyMetadataReader();
 
     /**
+     * Runs the specified command request.
+     */
+    virtual std::pair<rpc::UniqueReply, DBClientWithCommands*> runCommandWithTarget(
+        OpMsgRequest request);
+
+    /**
+     * Runs the specified command request. This thin wrapper just unwraps the reply and ignores the
+     * target connection from the above runCommandWithTarget().
+     */
+    rpc::UniqueReply runCommand(OpMsgRequest request) {
+        return runCommandWithTarget(std::move(request)).first;
+    }
+
+    /**
      * Runs a database command. This variant allows the caller to manually specify the metadata
      * for the request, and receive it for the reply.
      *
      * TODO: rename this to runCommand, and change the old one to runCommandLegacy.
      */
-    virtual rpc::UniqueReply runCommandWithMetadata(StringData database,
-                                                    StringData command,
-                                                    const BSONObj& metadata,
-                                                    const BSONObj& commandArgs);
+    rpc::UniqueReply runCommandWithMetadata(StringData database,
+                                            StringData command,
+                                            const BSONObj& metadata,
+                                            BSONObj commandArgs);
 
     /*
      * This wraps up the runCommandWithMetadata function above, but returns the DBClient that
@@ -427,11 +442,8 @@ public:
      *
      * This is used in the shell so that cursors can send getMore through the correct connection.
      */
-    virtual std::tuple<rpc::UniqueReply, DBClientWithCommands*> runCommandWithMetadataAndTarget(
-        StringData database,
-        StringData command,
-        const BSONObj& metadata,
-        const BSONObj& commandArgs);
+    std::tuple<rpc::UniqueReply, DBClientWithCommands*> runCommandWithMetadataAndTarget(
+        StringData database, StringData command, const BSONObj& metadata, BSONObj commandArgs);
 
     /** Run a database command.  Database commands are represented as BSON objects.  Common database
         commands have prebuilt helper functions -- see below.  If a helper is not available you can
@@ -446,10 +458,7 @@ public:
 
         @return true if the command returned "ok".
     */
-    virtual bool runCommand(const std::string& dbname,
-                            const BSONObj& cmd,
-                            BSONObj& info,
-                            int options = 0);
+    bool runCommand(const std::string& dbname, BSONObj cmd, BSONObj& info, int options = 0);
 
     /*
      * This wraps up the runCommand function avove, but returns the DBClient that actually ran
@@ -458,10 +467,10 @@ public:
      *
      * This is used in the shell so that cursors can send getMore through the correct connection.
      */
-    virtual std::tuple<bool, DBClientWithCommands*> runCommandWithTarget(const std::string& dbname,
-                                                                         const BSONObj& cmd,
-                                                                         BSONObj& info,
-                                                                         int options = 0);
+    std::tuple<bool, DBClientWithCommands*> runCommandWithTarget(const std::string& dbname,
+                                                                 BSONObj cmd,
+                                                                 BSONObj& info,
+                                                                 int options = 0);
 
     /**
     * Authenticates to another cluster member using appropriate authentication data.
@@ -1056,10 +1065,9 @@ public:
                                      const BSONObj* fieldsToReturn,
                                      int queryOptions);
 
-    virtual bool runCommand(const std::string& dbname,
-                            const BSONObj& cmd,
-                            BSONObj& info,
-                            int options = 0);
+    using DBClientWithCommands::runCommandWithTarget;
+    std::pair<rpc::UniqueReply, DBClientWithCommands*> runCommandWithTarget(
+        OpMsgRequest request) override;
 
     /**
        @return true if this connection is currently in a failed state.  When autoreconnect is on,
diff --git a/src/mongo/db/dbdirectclient.cpp b/src/mongo/db/dbdirectclient.cpp
index f7b71b88668..b85be86dd0d 100644
--- a/src/mongo/db/dbdirectclient.cpp
+++ b/src/mongo/db/dbdirectclient.cpp
@@ -70,7 +70,9 @@ private:
 }  // namespace
 
 
-DBDirectClient::DBDirectClient(OperationContext* opCtx) : _opCtx(opCtx) {}
+DBDirectClient::DBDirectClient(OperationContext* opCtx) : _opCtx(opCtx) {
+    _setServerRPCProtocols(rpc::supports::kAll);
+}
 
 bool DBDirectClient::isFailed() const {
     return false;
diff --git a/src/mongo/dbtests/commandtests.cpp b/src/mongo/dbtests/commandtests.cpp
index 7ad87efa08b..be9cd1ecf60 100644
--- a/src/mongo/dbtests/commandtests.cpp
+++ b/src/mongo/dbtests/commandtests.cpp
@@ -62,15 +62,9 @@ TEST(CommandTests, InputDocumentSequeceWorksEndToEnd) {
                               BSON("_id" << 5),
                           }}};
 
-    // Until the egress layer provides an API for document sequences, this needs to be handled
-    // manually.
-    Message rawReqest = request.serialize();
-    Message rawReply;
-    db.call(rawReqest, rawReply);
-    const auto reply = OpMsg::parse(rawReply);
-
-    ASSERT_BSONOBJ_EQ(reply.body, BSON("n" << 5 << "ok" << 1.0));
-    ASSERT(reply.sequences.empty());
+    const auto reply = db.runCommand(std::move(request));
+    ASSERT_EQ(int(reply->getProtocol()), int(rpc::Protocol::kOpMsg));
+    ASSERT_BSONOBJ_EQ(reply->getCommandReply(), BSON("n" << 5 << "ok" << 1.0));
     ASSERT_EQ(db.count(ns.ns()), 5u);
 }
 
diff --git a/src/mongo/dbtests/mock/mock_dbclient_connection.cpp b/src/mongo/dbtests/mock/mock_dbclient_connection.cpp
index 8fbcb3319b5..9e7f4467964 100644
--- a/src/mongo/dbtests/mock/mock_dbclient_connection.cpp
+++ b/src/mongo/dbtests/mock/mock_dbclient_connection.cpp
@@ -58,39 +58,20 @@ bool MockDBClientConnection::connect(const char* hostName,
     return false;
 }
 
-bool MockDBClientConnection::runCommand(const string& dbname,
-                                        const BSONObj& cmdObj,
-                                        BSONObj& info,
-                                        int options) {
-    checkConnection();
-
-    try {
-        return _remoteServer->runCommand(_remoteServerInstanceID, dbname, cmdObj, info, options);
-    } catch (const mongo::SocketException&) {
-        _isFailed = true;
-        throw;
-    }
-
-    return false;
-}
+std::pair<rpc::UniqueReply, DBClientWithCommands*> MockDBClientConnection::runCommandWithTarget(
+    OpMsgRequest request) {
 
-rpc::UniqueReply MockDBClientConnection::runCommandWithMetadata(StringData database,
-                                                                StringData command,
-                                                                const BSONObj& metadata,
-                                                                const BSONObj& commandArgs) {
     checkConnection();
 
     try {
-        return _remoteServer->runCommandWithMetadata(
-            _remoteServerInstanceID, database, command, metadata, commandArgs);
+        return {_remoteServer->runCommand(_remoteServerInstanceID, request), this};
     } catch (const mongo::SocketException&) {
         _isFailed = true;
         throw;
     }
-
-    MONGO_UNREACHABLE;
 }
 
+
 std::unique_ptr<mongo::DBClientCursor> MockDBClientConnection::query(const string& ns,
                                                                      mongo::Query query,
                                                                      int nToReturn,
diff --git a/src/mongo/dbtests/mock/mock_dbclient_connection.h b/src/mongo/dbtests/mock/mock_dbclient_connection.h
index d680fbdad99..bcb57be8aaa 100644
--- a/src/mongo/dbtests/mock/mock_dbclient_connection.h
+++ b/src/mongo/dbtests/mock/mock_dbclient_connection.h
@@ -64,15 +64,9 @@ public:
         return connect(host.toString().c_str(), applicationName, errmsg);
     }
 
-    bool runCommand(const std::string& dbname,
-                    const mongo::BSONObj& cmdObj,
-                    mongo::BSONObj& info,
-                    int options = 0);
-
-    rpc::UniqueReply runCommandWithMetadata(StringData database,
-                                            StringData command,
-                                            const BSONObj& metadata,
-                                            const BSONObj& commandArgs) final;
+    using DBClientWithCommands::runCommandWithTarget;
+    std::pair<rpc::UniqueReply, DBClientWithCommands*> runCommandWithTarget(
+        OpMsgRequest request) override;
 
     std::unique_ptr<mongo::DBClientCursor> query(const std::string& ns,
                                                  mongo::Query query = mongo::Query(),
diff --git a/src/mongo/dbtests/mock/mock_remote_db_server.cpp b/src/mongo/dbtests/mock/mock_remote_db_server.cpp
index 6851d4232ba..cf1380e99cb 100644
--- a/src/mongo/dbtests/mock/mock_remote_db_server.cpp
+++ b/src/mongo/dbtests/mock/mock_remote_db_server.cpp
@@ -136,20 +136,16 @@ void MockRemoteDBServer::remove(const string& ns, Query query, int flags) {
     _dataMgr.erase(ns);
 }
 
-rpc::UniqueReply MockRemoteDBServer::runCommandWithMetadata(MockRemoteDBServer::InstanceID id,
-                                                            StringData database,
-                                                            StringData commandName,
-                                                            const BSONObj& metadata,
-                                                            const BSONObj& commandArgs) {
+rpc::UniqueReply MockRemoteDBServer::runCommand(InstanceID id, const OpMsgRequest& request) {
     checkIfUp(id);
-    std::string cmdName = commandName.toString();
+    std::string cmdName = request.getCommandName().toString();
 
     BSONObj reply;
     {
         scoped_spinlock lk(_lock);
 
         uassert(ErrorCodes::IllegalOperation,
-                str::stream() << "no reply for command: " << commandName,
+                str::stream() << "no reply for command: " << cmdName,
                 _cmdMap.count(cmdName));
 
         reply = _cmdMap[cmdName]->next();
@@ -176,25 +172,6 @@ rpc::UniqueReply MockRemoteDBServer::runCommandWithMetadata(MockRemoteDBServer::
     return rpc::UniqueReply(std::move(message), std::move(replyView));
 }
 
-bool MockRemoteDBServer::runCommand(MockRemoteDBServer::InstanceID id,
-                                    const string& dbname,
-                                    const BSONObj& cmdObj,
-                                    BSONObj& info,
-                                    int options) {
-    BSONObj upconvertedRequest;
-    BSONObj upconvertedMetadata;
-    std::tie(upconvertedRequest, upconvertedMetadata) =
-        rpc::upconvertRequestMetadata(cmdObj, options);
-
-    StringData commandName = upconvertedRequest.firstElementFieldName();
-
-    auto res =
-        runCommandWithMetadata(id, dbname, commandName, upconvertedMetadata, upconvertedRequest);
-
-    info = res->getCommandReply().getOwned();
-    return info["ok"].trueValue();
-}
-
 mongo::BSONArray MockRemoteDBServer::query(MockRemoteDBServer::InstanceID id,
                                            const string& ns,
                                            mongo::Query query,
diff --git a/src/mongo/dbtests/mock/mock_remote_db_server.h b/src/mongo/dbtests/mock/mock_remote_db_server.h
index 791009ec459..743c26a3083 100644
--- a/src/mongo/dbtests/mock/mock_remote_db_server.h
+++ b/src/mongo/dbtests/mock/mock_remote_db_server.h
@@ -148,17 +148,7 @@ public:
     //
     // DBClientBase methods
     //
-    bool runCommand(InstanceID id,
-                    const std::string& dbname,
-                    const mongo::BSONObj& cmdObj,
-                    mongo::BSONObj& info,
-                    int options = 0);
-
-    rpc::UniqueReply runCommandWithMetadata(InstanceID id,
-                                            StringData database,
-                                            StringData commandName,
-                                            const BSONObj& metadata,
-                                            const BSONObj& commandArgs);
+    rpc::UniqueReply runCommand(InstanceID id, const OpMsgRequest& request);
 
     mongo::BSONArray query(InstanceID id,
                            const std::string& ns,
diff --git a/src/mongo/dbtests/mock_replica_set_test.cpp b/src/mongo/dbtests/mock_replica_set_test.cpp
index 8256cfaadc6..b7e7f2427b5 100644
--- a/src/mongo/dbtests/mock_replica_set_test.cpp
+++ b/src/mongo/dbtests/mock_replica_set_test.cpp
@@ -28,6 +28,7 @@
 #include "mongo/platform/basic.h"
 
 #include "mongo/client/dbclientinterface.h"
+#include "mongo/dbtests/mock/mock_dbclient_connection.h"
 #include "mongo/dbtests/mock/mock_replica_set.h"
 #include "mongo/unittest/unittest.h"
 
@@ -41,6 +42,7 @@ using mongo::BSONObjBuilder;
 using mongo::BSONObjIterator;
 using mongo::ConnectionString;
 using mongo::HostAndPort;
+using mongo::MockDBClientConnection;
 using mongo::MockRemoteDBServer;
 using mongo::MockReplicaSet;
 using mongo::repl::ReplSetConfig;
@@ -77,8 +79,8 @@ TEST(MockReplicaSetTest, IsMasterNode0) {
 
     BSONObj cmdResponse;
     MockRemoteDBServer* node = replSet.getNode("$n0:27017");
-    const MockRemoteDBServer::InstanceID id = node->getInstanceID();
-    bool ok = node->runCommand(id, "foo.bar", BSON("ismaster" << 1), cmdResponse);
+    bool ok =
+        MockDBClientConnection(node).runCommand("foo.bar", BSON("ismaster" << 1), cmdResponse);
     ASSERT(ok);
 
     ASSERT(cmdResponse["ismaster"].trueValue());
@@ -105,8 +107,8 @@ TEST(MockReplicaSetTest, IsMasterNode1) {
 
     BSONObj cmdResponse;
     MockRemoteDBServer* node = replSet.getNode("$n1:27017");
-    const MockRemoteDBServer::InstanceID id = node->getInstanceID();
-    bool ok = node->runCommand(id, "foo.bar", BSON("ismaster" << 1), cmdResponse);
+    bool ok =
+        MockDBClientConnection(node).runCommand("foo.bar", BSON("ismaster" << 1), cmdResponse);
     ASSERT(ok);
 
     ASSERT(!cmdResponse["ismaster"].trueValue());
@@ -133,8 +135,8 @@ TEST(MockReplicaSetTest, IsMasterNode2) {
 
     BSONObj cmdResponse;
     MockRemoteDBServer* node = replSet.getNode("$n2:27017");
-    const MockRemoteDBServer::InstanceID id = node->getInstanceID();
-    bool ok = node->runCommand(id, "foo.bar", BSON("ismaster" << 1), cmdResponse);
+    bool ok =
+        MockDBClientConnection(node).runCommand("foo.bar", BSON("ismaster" << 1), cmdResponse);
     ASSERT(ok);
 
     ASSERT(!cmdResponse["ismaster"].trueValue());
@@ -161,8 +163,8 @@ TEST(MockReplicaSetTest, ReplSetGetStatusNode0) {
 
     BSONObj cmdResponse;
     MockRemoteDBServer* node = replSet.getNode("$n0:27017");
-    const MockRemoteDBServer::InstanceID id = node->getInstanceID();
-    bool ok = node->runCommand(id, "foo.bar", BSON("replSetGetStatus" << 1), cmdResponse);
+    bool ok = MockDBClientConnection(node).runCommand(
+        "foo.bar", BSON("replSetGetStatus" << 1), cmdResponse);
     ASSERT(ok);
 
     ASSERT_EQUALS("n", cmdResponse["set"].str());
@@ -194,8 +196,8 @@ TEST(MockReplicaSetTest, ReplSetGetStatusNode1) {
 
     BSONObj cmdResponse;
     MockRemoteDBServer* node = replSet.getNode("$n1:27017");
-    const MockRemoteDBServer::InstanceID id = node->getInstanceID();
-    bool ok = node->runCommand(id, "foo.bar", BSON("replSetGetStatus" << 1), cmdResponse);
+    bool ok = MockDBClientConnection(node).runCommand(
+        "foo.bar", BSON("replSetGetStatus" << 1), cmdResponse);
     ASSERT(ok);
 
     ASSERT_EQUALS("n", cmdResponse["set"].str());
@@ -229,8 +231,8 @@ TEST(MockReplicaSetTest, ReplSetGetStatusNode2) {
 
     BSONObj cmdResponse;
     MockRemoteDBServer* node = replSet.getNode("$n2:27017");
-    const MockRemoteDBServer::InstanceID id = node->getInstanceID();
-    bool ok = node->runCommand(id, "foo.bar", BSON("replSetGetStatus" << 1), cmdResponse);
+    bool ok = MockDBClientConnection(node).runCommand(
+        "foo.bar", BSON("replSetGetStatus" << 1), cmdResponse);
     ASSERT(ok);
 
     ASSERT_EQUALS("n", cmdResponse["set"].str());
@@ -299,8 +301,8 @@ TEST(MockReplicaSetTest, IsMasterReconfigNodeRemoved) {
         // Check isMaster for node still in set
         BSONObj cmdResponse;
         MockRemoteDBServer* node = replSet.getNode("$n0:27017");
-        const MockRemoteDBServer::InstanceID id = node->getInstanceID();
-        bool ok = node->runCommand(id, "foo.bar", BSON("ismaster" << 1), cmdResponse);
+        bool ok =
+            MockDBClientConnection(node).runCommand("foo.bar", BSON("ismaster" << 1), cmdResponse);
         ASSERT(ok);
 
         ASSERT(cmdResponse["ismaster"].trueValue());
@@ -327,8 +329,8 @@ TEST(MockReplicaSetTest, IsMasterReconfigNodeRemoved) {
         // Check isMaster for node still not in set anymore
         BSONObj cmdResponse;
         MockRemoteDBServer* node = replSet.getNode(hostToRemove);
-        const MockRemoteDBServer::InstanceID id = node->getInstanceID();
-        bool ok = node->runCommand(id, "foo.bar", BSON("ismaster" << 1), cmdResponse);
+        bool ok =
+            MockDBClientConnection(node).runCommand("foo.bar", BSON("ismaster" << 1), cmdResponse);
         ASSERT(ok);
 
         ASSERT(!cmdResponse["ismaster"].trueValue());
@@ -350,8 +352,8 @@ TEST(MockReplicaSetTest, replSetGetStatusReconfigNodeRemoved) {
         // Check replSetGetStatus for node still in set
         BSONObj cmdResponse;
         MockRemoteDBServer* node = replSet.getNode("$n2:27017");
-        const MockRemoteDBServer::InstanceID id = node->getInstanceID();
-        bool ok = node->runCommand(id, "foo.bar", BSON("replSetGetStatus" << 1), cmdResponse);
+        bool ok = MockDBClientConnection(node).runCommand(
+            "foo.bar", BSON("replSetGetStatus" << 1), cmdResponse);
         ASSERT(ok);
 
         ASSERT_EQUALS("n", cmdResponse["set"].str());
@@ -383,8 +385,8 @@ TEST(MockReplicaSetTest, replSetGetStatusReconfigNodeRemoved) {
         // Check replSetGetStatus for node still not in set anymore
         BSONObj cmdResponse;
         MockRemoteDBServer* node = replSet.getNode(hostToRemove);
-        const MockRemoteDBServer::InstanceID id = node->getInstanceID();
-        bool ok = node->runCommand(id, "foo.bar", BSON("replSetGetStatus" << 1), cmdResponse);
+        bool ok = MockDBClientConnection(node).runCommand(
+            "foo.bar", BSON("replSetGetStatus" << 1), cmdResponse);
         ASSERT(ok);
 
         ASSERT_EQUALS("n", cmdResponse["set"].str());