SERVER-73189 Change DbClient runCommandWithTarget to take in a DatabaseName object

3 files changed, 41 insertions, 48 deletions

diff --git a/src/mongo/shell/bench.cpp b/src/mongo/shell/bench.cpp
index f8f4a29b26d..7e4ee3ae10d 100644
--- a/src/mongo/shell/bench.cpp
+++ b/src/mongo/shell/bench.cpp
@@ -118,7 +118,7 @@ BSONObj fixQuery(const BSONObj& obj, BsonTemplateEvaluator& btl) {
 }
 
 bool runCommandWithSession(DBClientBase* conn,
-                           const std::string& dbname,
+                           const DatabaseName& dbName,
                            const BSONObj& cmdObj,
                            int options,
                            const boost::optional<LogicalSessionIdToClient>& lsid,
@@ -126,8 +126,7 @@ bool runCommandWithSession(DBClientBase* conn,
                            BSONObj* result) {
     if (!lsid) {
         invariant(!txnNumber);
-        // Shell is not tenant aware, so use boost::none here.
-        return conn->runCommand(DatabaseName(boost::none, dbname), cmdObj, *result);
+        return conn->runCommand(dbName, cmdObj, *result);
     }
 
     BSONObjBuilder cmdObjWithLsidBuilder;
@@ -162,18 +161,16 @@ bool runCommandWithSession(DBClientBase* conn,
         cmdObjWithLsidBuilder.append("startTransaction", true);
     }
 
-    // Shell is not tenant aware, so use boost::none here.
-    return conn->runCommand(
-        DatabaseName(boost::none, dbname), cmdObjWithLsidBuilder.done(), *result);
+    return conn->runCommand(dbName, cmdObjWithLsidBuilder.done(), *result);
 }
 
 bool runCommandWithSession(DBClientBase* conn,
-                           const std::string& dbname,
+                           const DatabaseName& dbName,
                            const BSONObj& cmdObj,
                            int options,
                            const boost::optional<LogicalSessionIdToClient>& lsid,
                            BSONObj* result) {
-    return runCommandWithSession(conn, dbname, cmdObj, options, lsid, boost::none, result);
+    return runCommandWithSession(conn, dbName, cmdObj, options, lsid, boost::none, result);
 }
 
 void abortTransaction(DBClientBase* conn,
@@ -182,7 +179,7 @@ void abortTransaction(DBClientBase* conn,
     BSONObj abortTransactionCmd = BSON("abortTransaction" << 1);
     BSONObj abortCommandResult;
     const bool successful = runCommandWithSession(conn,
-                                                  "admin",
+                                                  DatabaseName(boost::none, "admin"),
                                                   abortTransactionCmd,
                                                   kMultiStatementTransactionOption,
                                                   lsid,
@@ -212,11 +209,8 @@ int runQueryWithReadCommands(DBClientBase* conn,
                              Milliseconds delayBeforeGetMore,
                              BSONObj readPrefObj,
                              BSONObj* objOut) {
-    const auto dbName = findCommand->getNamespaceOrUUID()
-                            .nss()
-                            .value_or(NamespaceString())
-                            .dbName()
-                            .toStringWithTenantId();
+    const auto dbName =
+        findCommand->getNamespaceOrUUID().nss().value_or(NamespaceString()).dbName();
 
     BSONObj findCommandResult;
     BSONObj findCommandObj = findCommand->toBSON(readPrefObj);
@@ -286,8 +280,6 @@ Timestamp getLatestClusterTime(DBClientBase* conn) {
     findCommand->setLimit(1LL);
     findCommand->setSingleBatch(true);
     invariant(query_request_helper::validateFindCommandRequest(*findCommand));
-    const auto dbName =
-        findCommand->getNamespaceOrUUID().nss().value_or(NamespaceString()).db().toString();
 
     BSONObj oplogResult;
     int count = runQueryWithReadCommands(conn,
@@ -799,7 +791,9 @@ std::unique_ptr<DBClientBase> BenchRunConfig::createConnectionImpl(
 }
 
 std::unique_ptr<DBClientBase> BenchRunConfig::createConnection() const {
-    return BenchRunConfig::createConnectionImpl(*this);
+    auto conn = BenchRunConfig::createConnectionImpl(*this);
+    conn->setAlwaysAppendDollarTenant_forTest();
+    return conn;
 }
 
 BenchRunState::BenchRunState(unsigned numWorkers)
@@ -1055,13 +1049,14 @@ void BenchRunOp::executeOnce(DBClientBase* conn,
         case OpType::FINDONE: {
             BSONObj fixedQuery = fixQuery(this->query, *state->bsonTemplateEvaluator);
             BSONObj result;
-            auto findCommand = std::make_unique<FindCommandRequest>(NamespaceString(this->ns));
+            auto findCommand =
+                std::make_unique<FindCommandRequest>(NamespaceString(this->tenantId, this->ns));
             findCommand->setFilter(fixedQuery);
             findCommand->setProjection(this->projection);
             findCommand->setLimit(1LL);
             findCommand->setSingleBatch(true);
             findCommand->setSort(this->sort);
-            findCommand->setDollarTenant(this->tenantId);
+
             if (config.useSnapshotReads) {
                 findCommand->setReadConcern(readConcernSnapshot);
             }
@@ -1099,7 +1094,7 @@ void BenchRunOp::executeOnce(DBClientBase* conn,
             {
                 BenchRunEventTrace _bret(&state->stats->commandCounter);
                 ok = runCommandWithSession(conn,
-                                           this->ns,
+                                           DatabaseName(this->tenantId, this->ns),
                                            fixQuery(this->command, *state->bsonTemplateEvaluator),
                                            this->options,
                                            lsid,
@@ -1121,7 +1116,7 @@ void BenchRunOp::executeOnce(DBClientBase* conn,
                             str::stream()
                                 << "getMore command failed; reply was: " << getMoreCommandResult,
                             runCommandWithSession(conn,
-                                                  this->ns,
+                                                  DatabaseName(this->tenantId, this->ns),
                                                   getMoreRequest.toBSON({}),
                                                   kNoOptions,
                                                   lsid,
@@ -1142,7 +1137,8 @@ void BenchRunOp::executeOnce(DBClientBase* conn,
             uassert(
                 28824, "cannot use 'options' in combination with read commands", !this->options);
 
-            auto findCommand = std::make_unique<FindCommandRequest>(NamespaceString(this->ns));
+            auto findCommand =
+                std::make_unique<FindCommandRequest>(NamespaceString(this->tenantId, this->ns));
             findCommand->setFilter(fixedQuery);
             findCommand->setProjection(this->projection);
             if (this->skip) {
@@ -1157,7 +1153,6 @@ void BenchRunOp::executeOnce(DBClientBase* conn,
             if (!this->sort.isEmpty()) {
                 findCommand->setSort(this->sort);
             }
-            findCommand->setDollarTenant(this->tenantId);
 
             BSONObjBuilder readConcernBuilder;
             if (config.useSnapshotReads) {
@@ -1214,9 +1209,6 @@ void BenchRunOp::executeOnce(DBClientBase* conn,
 
                 BSONObjBuilder builder;
                 builder.append("update", nsToCollectionSubstring(this->ns));
-                if (this->tenantId) {
-                    this->tenantId->serializeToBSON("$tenant"_sd, &builder);
-                }
 
                 BSONArrayBuilder updateArray(builder.subarrayStart("updates"));
                 {
@@ -1266,7 +1258,7 @@ void BenchRunOp::executeOnce(DBClientBase* conn,
                     txnNumberForOp = state->txnNumber;
                 }
                 runCommandWithSession(conn,
-                                      nsToDatabaseSubstring(this->ns).toString(),
+                                      DatabaseName(this->tenantId, nsToDatabaseSubstring(this->ns)),
                                       builder.done(),
                                       kNoOptions,
                                       lsid,
@@ -1284,9 +1276,6 @@ void BenchRunOp::executeOnce(DBClientBase* conn,
                 BSONObj insertDoc;
                 BSONObjBuilder builder;
                 builder.append("insert", nsToCollectionSubstring(this->ns));
-                if (this->tenantId) {
-                    this->tenantId->serializeToBSON("$tenant"_sd, &builder);
-                }
 
                 BSONArrayBuilder docBuilder(builder.subarrayStart("documents"));
                 if (this->isDocAnArray) {
@@ -1307,7 +1296,7 @@ void BenchRunOp::executeOnce(DBClientBase* conn,
                     txnNumberForOp = state->txnNumber;
                 }
                 runCommandWithSession(conn,
-                                      nsToDatabaseSubstring(this->ns).toString(),
+                                      DatabaseName(this->tenantId, nsToDatabaseSubstring(this->ns)),
                                       builder.done(),
                                       kNoOptions,
                                       lsid,
@@ -1324,9 +1313,6 @@ void BenchRunOp::executeOnce(DBClientBase* conn,
                 BSONObj predicate = fixQuery(this->query, *state->bsonTemplateEvaluator);
                 BSONObjBuilder builder;
                 builder.append("delete", nsToCollectionSubstring(this->ns));
-                if (this->tenantId) {
-                    this->tenantId->serializeToBSON("$tenant"_sd, &builder);
-                }
 
                 BSONArrayBuilder docBuilder(builder.subarrayStart("deletes"));
                 int limit = (this->multi == true) ? 0 : 1;
@@ -1340,7 +1326,7 @@ void BenchRunOp::executeOnce(DBClientBase* conn,
                     txnNumberForOp = state->txnNumber;
                 }
                 runCommandWithSession(conn,
-                                      nsToDatabaseSubstring(this->ns).toString(),
+                                      DatabaseName(this->tenantId, nsToDatabaseSubstring(this->ns)),
                                       builder.done(),
                                       kNoOptions,
                                       lsid,
@@ -1353,14 +1339,12 @@ void BenchRunOp::executeOnce(DBClientBase* conn,
         case OpType::CREATEINDEX:
             conn->createIndex(NamespaceString(this->tenantId, this->ns),
                               this->key,
-                              boost::none /* writeConcernObj */,
-                              true);
+                              boost::none /* writeConcernObj */);
             break;
         case OpType::DROPINDEX:
             conn->dropIndex(NamespaceString(this->tenantId, this->ns),
                             this->key,
-                            boost::none /* writeConcernObj */,
-                            true);
+                            boost::none /* writeConcernObj */);
             break;
         case OpType::LET: {
             BSONObjBuilder templateBuilder;
diff --git a/src/mongo/shell/encrypted_dbclient_base.cpp b/src/mongo/shell/encrypted_dbclient_base.cpp
index fb2e3b90d09..7758ed71f98 100644
--- a/src/mongo/shell/encrypted_dbclient_base.cpp
+++ b/src/mongo/shell/encrypted_dbclient_base.cpp
@@ -97,6 +97,7 @@ EncryptedDBClientBase::EncryptedDBClientBase(std::unique_ptr<DBClientBase> conn,
     : _conn(std::move(conn)), _encryptionOptions(std::move(encryptionOptions)), _cx(cx) {
     validateCollection(cx, collection);
     _collection = JS::Heap<JS::Value>(collection);
+    _conn->setAlwaysAppendDollarTenant_forTest();
 };
 
 std::string EncryptedDBClientBase::getServerAddress() const {
@@ -113,7 +114,7 @@ void EncryptedDBClientBase::say(Message& toSend, bool isRetry, std::string* actu
 
 BSONObj EncryptedDBClientBase::encryptDecryptCommand(const BSONObj& object,
                                                      bool encrypt,
-                                                     const StringData databaseName) {
+                                                     const DatabaseName& dbName) {
     std::stack<std::pair<BSONObjIterator, BSONObjBuilder>> frameStack;
 
     // The encryptDecryptCommand frameStack requires a guard because  if encryptMarking or
@@ -164,7 +165,12 @@ BSONObj EncryptedDBClientBase::encryptDecryptCommand(const BSONObj& object,
         }
     }
     invariant(frameStack.size() == 1);
-    frameStack.top().second.append("$db", databaseName);
+    // Append '$db' which shouldn't contain tenantid.
+    frameStack.top().second.append("$db", dbName.toString());
+    // If encrypt request, append '$tenant' which contains tenantid.
+    if (encrypt && dbName.tenantId() && !object.hasField("$tenant")) {
+        dbName.tenantId()->serializeToBSON("$tenant", &frameStack.top().second);
+    }
     return frameStack.top().second.obj();
 }
 
@@ -194,9 +200,9 @@ void EncryptedDBClientBase::decryptPayload(ConstDataRange data,
 }
 
 EncryptedDBClientBase::RunCommandReturn EncryptedDBClientBase::processResponseFLE1(
-    EncryptedDBClientBase::RunCommandReturn result, const StringData databaseName) {
+    EncryptedDBClientBase::RunCommandReturn result, const DatabaseName& dbName) {
     auto rawReply = result.returnReply->getCommandReply();
-    return prepareReply(std::move(result), encryptDecryptCommand(rawReply, false, databaseName));
+    return prepareReply(std::move(result), encryptDecryptCommand(rawReply, false, dbName));
 }
 
 EncryptedDBClientBase::RunCommandReturn EncryptedDBClientBase::processResponseFLE2(
@@ -230,8 +236,11 @@ EncryptedDBClientBase::RunCommandReturn EncryptedDBClientBase::doRunCommand(
 
 EncryptedDBClientBase::RunCommandReturn EncryptedDBClientBase::handleEncryptionRequest(
     EncryptedDBClientBase::RunCommandParams params) {
-    auto commandName = params.request.getCommandName().toString();
-    auto databaseName = params.request.getDatabase().toString();
+    auto& request = params.request;
+    auto commandName = request.getCommandName().toString();
+    const DatabaseName dbName = request.body.hasField("$tenant")
+        ? DatabaseName(TenantId(request.body["$tenant"].OID()), request.getDatabase())
+        : DatabaseName(boost::none, request.getDatabase());
 
     if (std::find(kEncryptedCommands.begin(), kEncryptedCommands.end(), StringData(commandName)) ==
         std::end(kEncryptedCommands)) {
@@ -239,7 +248,7 @@ EncryptedDBClientBase::RunCommandReturn EncryptedDBClientBase::handleEncryptionR
     }
 
     EncryptedDBClientBase::RunCommandReturn result(doRunCommand(std::move(params)));
-    return processResponseFLE1(processResponseFLE2(std::move(result)), databaseName);
+    return processResponseFLE1(processResponseFLE2(std::move(result)), dbName);
 }
 
 std::pair<rpc::UniqueReply, DBClientBase*> EncryptedDBClientBase::runCommandWithTarget(
diff --git a/src/mongo/shell/encrypted_dbclient_base.h b/src/mongo/shell/encrypted_dbclient_base.h
index 897bec92474..73854b39a2b 100644
--- a/src/mongo/shell/encrypted_dbclient_base.h
+++ b/src/mongo/shell/encrypted_dbclient_base.h
@@ -200,13 +200,13 @@ protected:
 
     virtual RunCommandReturn handleEncryptionRequest(RunCommandParams params);
 
-    RunCommandReturn processResponseFLE1(RunCommandReturn result, StringData databaseName);
+    RunCommandReturn processResponseFLE1(RunCommandReturn result, const DatabaseName& databaseName);
 
     RunCommandReturn processResponseFLE2(RunCommandReturn result);
 
     RunCommandReturn prepareReply(RunCommandReturn result, BSONObj decryptedDoc);
 
-    BSONObj encryptDecryptCommand(const BSONObj& object, bool encrypt, StringData databaseName);
+    BSONObj encryptDecryptCommand(const BSONObj& object, bool encrypt, const DatabaseName& dbName);
 
     JS::Value getCollection() const;