SERVER-35758 Mongo shell prompt errors when running transactions after overriding 'db'

(cherry picked from commit bb9f225e47faacb95ae9e5658313a5ebf56f37bf)

3 files changed, 67 insertions, 1 deletions

diff --git a/jstests/core/txns/shell_prompt_in_transaction.js b/jstests/core/txns/shell_prompt_in_transaction.js
new file mode 100644
index 00000000000..b843a1b7271
--- /dev/null
+++ b/jstests/core/txns/shell_prompt_in_transaction.js
@@ -0,0 +1,43 @@
+// Test shell prompt doesn't run in the session of the global 'db'.
+// @tags: [uses_transactions]
+
+(function() {
+    "use strict";
+
+    const collName = "shell_prompt_in_transaction";
+
+    db.getCollection(collName).drop();
+    assert.commandWorked(db.runCommand({create: collName, writeConcern: {w: "majority"}}));
+
+    // Override the global "db".
+    const session = db.getMongo().startSession();
+    db = session.getDatabase(db.getName());
+    const coll = db.getCollection(collName);
+
+    function simulatePrompt() {
+        __promptWrapper__(defaultPrompt);
+    }
+
+    // Start a transaction, so the session will attach txn info to the commands running on it.
+    session.startTransaction();
+    jsTestLog("Run shell prompt to simulate a user hitting enter.");
+    simulatePrompt();
+    const doc = {_id: "shell-write"};
+    assert.commandWorked(coll.insert(doc));
+    assert.docEq(doc, coll.findOne());
+    simulatePrompt();
+    session.abortTransaction();
+    assert.docEq(null, coll.findOne());
+
+    // Start a transaction, so the session has a running transaction now.
+    simulatePrompt();
+    session.startTransaction();
+    jsTestLog("Run shell prompt to simulate a user hitting enter.");
+    simulatePrompt();
+    assert.commandWorked(coll.insert(doc));
+    simulatePrompt();
+    session.commitTransaction();
+    assert.docEq(doc, coll.findOne());
+
+    coll.drop();
+})();
diff --git a/src/mongo/shell/dbshell.cpp b/src/mongo/shell/dbshell.cpp
index 9fb501779f2..c2f2e3013f3 100644
--- a/src/mongo/shell/dbshell.cpp
+++ b/src/mongo/shell/dbshell.cpp
@@ -560,7 +560,7 @@ string finishCode(string code) {
 }
 
 bool execPrompt(mongo::Scope& scope, const char* promptFunction, string& prompt) {
-    string execStatement = string("__prompt__ = ") + promptFunction + "();";
+    string execStatement = string("__promptWrapper__(") + promptFunction + ");";
     scope.exec("delete __prompt__;", "", false, false, false, 0);
     scope.exec(execStatement, "", false, false, false, 0);
     if (scope.type("__prompt__") == String) {
diff --git a/src/mongo/shell/utils.js b/src/mongo/shell/utils.js
index 3820eac4f36..041eb0d6bf5 100644
--- a/src/mongo/shell/utils.js
+++ b/src/mongo/shell/utils.js
@@ -1044,6 +1044,29 @@ shellHelper.show = function(what) {
 
 };
 
+__promptWrapper__ = function(promptFunction) {
+    // Call promptFunction directly if the global "db" is not defined, e.g. --nodb.
+    if (typeof db === 'undefined' || !(db instanceof DB)) {
+        __prompt__ = promptFunction();
+        return;
+    }
+
+    // Stash the global "db" for the prompt function to make sure the session
+    // of the global "db" isn't accessed by the prompt function.
+    let originalDB = db;
+    try {
+        db = originalDB.getMongo().getDB(originalDB.getName());
+        // Setting db._session to be a _DummyDriverSession instance makes it so that
+        // a logical session id isn't included in the isMaster and replSetGetStatus
+        // commands and therefore won't interfere with the session associated with the
+        // global "db" object.
+        db._session = new _DummyDriverSession(db.getMongo());
+        __prompt__ = promptFunction();
+    } finally {
+        db = originalDB;
+    }
+};
+
 Math.sigFig = function(x, N) {
     if (!N) {
         N = 3;