From 59f8b1115bc749c545b7669bc2c324525f50cb94 Mon Sep 17 00:00:00 2001
From: Anna Henningsen <anna.henningsen@mongodb.com>
Date: Mon, 21 Nov 2022 14:55:37 +0100
Subject: [PATCH] feat(shell-api): accept another Mongo instance as 1st arg to
 ctor MONGOSH-1272

---
 packages/cli-repl/src/smoke-tests-fle.ts | 4 ++--
 packages/cli-repl/test/e2e-fle.spec.ts   | 4 ++--
 packages/shell-api/src/mongo.ts          | 6 ++++--
 3 files changed, 8 insertions(+), 6 deletions(-)

diff --git a/packages/cli-repl/src/smoke-tests-fle.ts b/packages/cli-repl/src/smoke-tests-fle.ts
index e2956773c7..e2c367e47f 100644
--- a/packages/cli-repl/src/smoke-tests-fle.ts
+++ b/packages/cli-repl/src/smoke-tests-fle.ts
@@ -26,7 +26,7 @@ assert(db.getName() === dbname, 'db name must match');
 
 const local = { key: Buffer.from('kh4Gv2N8qopZQMQYMEtww/AkPsIrXNmEMxTrs3tUoTQZbZu4msdRUaR8U5fXD7A7QXYHcEvuu4WctJLoT+NvvV3eeIg3MD+K8H9SR794m/safgRHdIfy6PD+rFpvmFbY', 'base64') };
 
-const keyMongo = Mongo(db.getMongo()._uri, {
+const keyMongo = Mongo(db.getMongo(), {
   keyVaultNamespace: dbname + '.__keyVault',
   kmsProviders: { local }
 });
@@ -51,7 +51,7 @@ schemaMap[dbname + '.employees'] = {
 
 console.log('Using schema map', schemaMap);
 
-const autoMongo = Mongo(db.getMongo()._uri, {
+const autoMongo = Mongo(db.getMongo(), {
   keyVaultNamespace: dbname + '.__keyVault',
   kmsProviders: { local },
   schemaMap
diff --git a/packages/cli-repl/test/e2e-fle.spec.ts b/packages/cli-repl/test/e2e-fle.spec.ts
index 1dd4ee58b9..65ab40d55d 100644
--- a/packages/cli-repl/test/e2e-fle.spec.ts
+++ b/packages/cli-repl/test/e2e-fle.spec.ts
@@ -437,7 +437,7 @@ describe('FLE tests', () => {
     // Wrapper for executeLine that expects single-line output
     const runSingleLine = async(line) => (await shell.executeLine(line)).split('\n')[0].trim();
     await runSingleLine('local = { key: BinData(0, "kh4Gv2N8qopZQMQYMEtww/AkPsIrXNmEMxTrs3tUoTQZbZu4msdRUaR8U5fXD7A7QXYHcEvuu4WctJLoT+NvvV3eeIg3MD+K8H9SR794m/safgRHdIfy6PD+rFpvmFbY") }');
-    await runSingleLine(`keyMongo = Mongo(db.getMongo()._uri, { \
+    await runSingleLine(`keyMongo = Mongo(db.getMongo(), { \
       keyVaultNamespace: '${dbname}.keyVault', \
       kmsProviders: { local }, \
       explicitEncryptionOnly: true \
@@ -473,7 +473,7 @@ describe('FLE tests', () => {
     // Wrapper for executeLine that expects single-line output
     const runSingleLine = async(line) => (await shell.executeLine(line)).split('\n')[0].trim();
     await runSingleLine('local = { key: BinData(0, "kh4Gv2N8qopZQMQYMEtww/AkPsIrXNmEMxTrs3tUoTQZbZu4msdRUaR8U5fXD7A7QXYHcEvuu4WctJLoT+NvvV3eeIg3MD+K8H9SR794m/safgRHdIfy6PD+rFpvmFbY") }');
-    await runSingleLine(`keyMongo = Mongo(db.getMongo()._uri, { \
+    await runSingleLine(`keyMongo = Mongo(db.getMongo(), { \
       keyVaultNamespace: '${dbname}.keyVault', \
       kmsProviders: { local }, \
       explicitEncryptionOnly: true \
diff --git a/packages/shell-api/src/mongo.ts b/packages/shell-api/src/mongo.ts
index 74f5838655..b4ba65224c 100644
--- a/packages/shell-api/src/mongo.ts
+++ b/packages/shell-api/src/mongo.ts
@@ -81,7 +81,7 @@ export default class Mongo extends ShellApiClass {
 
   constructor(
     instanceState: ShellInstanceState,
-    uri?: string,
+    uri?: string | Mongo,
     fleOptions?: ClientSideFieldLevelEncryptionOptions,
     otherOptions?: { api?: ServerApi | ServerApiVersion },
     sp?: ServiceProvider
@@ -92,7 +92,9 @@ export default class Mongo extends ShellApiClass {
     if (sp) {
       this.__serviceProvider = sp;
     }
-    if (typeof uri !== 'string') {
+    if (typeof uri === 'object' && uri !== null && typeof uri._uri === 'string') {
+      uri = uri._uri;
+    } else if (typeof uri !== 'string') {
       uri = sp?.getURI?.() ?? 'mongodb://localhost/';
     }
     this._connectionInfo = generateConnectionInfoFromCliArgs({