From 597e2bccb8b73b63549f7354077762f64858012a Mon Sep 17 00:00:00 2001
From: Anna Henningsen <anna.henningsen@mongodb.com>
Date: Wed, 21 Apr 2021 19:36:16 +0200
Subject: [PATCH] fix(shell-api): fix prompt for Atlas-proxy-style mongos
 MONGOSH-680
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

When connected to a server that we know is an Atlas server but
announces itself as a mongos, print “atlas proxy” in the prompt
rather than “mongos” to avoid confusion.

As a drive-by fix, make sure to use the service provider of the
current database, and not necessarily the initial one, as new
connections can be assigned to `db`.
---
 packages/cli-repl/src/mongosh-repl.spec.ts    | 31 +++++++++++++++++++
 packages/shell-api/src/mongo.ts               |  2 +-
 .../src/shell-internal-state.spec.ts          | 20 ++++++++++++
 .../shell-api/src/shell-internal-state.ts     | 18 ++++++++---
 4 files changed, 66 insertions(+), 5 deletions(-)

diff --git a/packages/cli-repl/src/mongosh-repl.spec.ts b/packages/cli-repl/src/mongosh-repl.spec.ts
index 74b8cdb635..4f79124db2 100644
--- a/packages/cli-repl/src/mongosh-repl.spec.ts
+++ b/packages/cli-repl/src/mongosh-repl.spec.ts
@@ -643,5 +643,36 @@ describe('MongoshNodeRepl', () => {
       expect(output).to.contain('> ');
       expect(output).to.not.contain('error');
     });
+
+    it('changes the prompt when db is reassigned', async() => {
+      const connectionInfo = {
+        extraInfo: {
+          uri: 'mongodb://localhost:27017/test',
+          is_localhost: true
+        },
+        buildInfo: {
+          version: '4.4.1',
+          modules: ['enterprise']
+        }
+      };
+
+      sp.getConnectionInfo.resolves(connectionInfo);
+      sp.getNewConnection.callsFake(async() => {
+        Object.assign(connectionInfo.extraInfo, {
+          is_localhost: true,
+          is_data_lake: true
+        });
+        return sp;
+      });
+      sp.platform = 2; // ReplPlatform.CLI ... let's maybe stop using an enum for this
+
+      const initialized = await mongoshRepl.initialize(serviceProvider);
+      await mongoshRepl.startRepl(initialized);
+      expect(output).to.contain('Enterprise > ');
+
+      input.write('db = Mongo("foo").getDB("bar")\n');
+      await waitEval(bus);
+      expect(output).to.contain('Atlas Data Lake > ');
+    });
   });
 });
diff --git a/packages/shell-api/src/mongo.ts b/packages/shell-api/src/mongo.ts
index 752bb2951f..3f51ae63ed 100644
--- a/packages/shell-api/src/mongo.ts
+++ b/packages/shell-api/src/mongo.ts
@@ -117,7 +117,7 @@ export default class Mongo extends ShellApiClass {
   // if used too early.
   get _serviceProvider(): ServiceProvider {
     if (this.__serviceProvider === null) {
-      throw new MongoshInternalError('No ServiceProvider available for this mongo');
+      throw new MongoshInternalError('No ServiceProvider available for this mongo', ShellApiErrors.NotConnected);
     }
     return this.__serviceProvider;
   }
diff --git a/packages/shell-api/src/shell-internal-state.spec.ts b/packages/shell-api/src/shell-internal-state.spec.ts
index d029586884..dca8fffb42 100644
--- a/packages/shell-api/src/shell-internal-state.spec.ts
+++ b/packages/shell-api/src/shell-internal-state.spec.ts
@@ -249,6 +249,26 @@ describe('ShellInternalState', () => {
       });
     });
 
+    describe('topology Sharded but it’s Atlas', () => {
+      it('shows atlas proxy identifier', async() => {
+        serviceProvider.getTopology.returns({
+          description: {
+            type: 'Sharded'
+          }
+        });
+        serviceProvider.getConnectionInfo.resolves({
+          extraInfo: {
+            uri: 'mongodb://localhost/',
+            is_atlas: true
+          }
+        });
+
+        await internalState.fetchConnectionInfo();
+        const prompt = await internalState.getDefaultPrompt();
+        expect(prompt).to.equal('[atlas proxy]> ');
+      });
+    });
+
     describe('topology Unknown', () => {
       it('just shows the default prompt', async() => {
         const servers = new Map();
diff --git a/packages/shell-api/src/shell-internal-state.ts b/packages/shell-api/src/shell-internal-state.ts
index 71266676b2..7eaa71e1e5 100644
--- a/packages/shell-api/src/shell-internal-state.ts
+++ b/packages/shell-api/src/shell-internal-state.ts
@@ -242,6 +242,17 @@ export default class ShellInternalState {
     );
   }
 
+  get currentServiceProvider(): ServiceProvider {
+    try {
+      return this.currentDb._mongo._serviceProvider;
+    } catch (err) {
+      if (err.code === ShellApiErrors.NotConnected) {
+        return this.initialServiceProvider;
+      }
+      throw err;
+    }
+  }
+
   public emitApiCall(event: ApiEvent): void {
     this.messageBus.emit('mongosh:api-call', event);
   }
@@ -255,7 +266,7 @@ export default class ShellInternalState {
       // eslint-disable-next-line complexity
       topology: () => {
         let topology: Topologies;
-        const topologyDescription = this.initialServiceProvider.getTopology()?.description as TopologyDescription;
+        const topologyDescription = this.currentServiceProvider.getTopology()?.description as TopologyDescription;
         const topologyType: TopologyTypeId | undefined = topologyDescription?.type;
         switch (topologyType) {
           case 'ReplicaSetNoPrimary':
@@ -326,12 +337,11 @@ export default class ShellInternalState {
   }
 
   private getTopologySpecificPrompt(): string {
-    const description = this.initialServiceProvider.getTopology()?.description;
+    const description = this.currentServiceProvider.getTopology()?.description;
     if (!description) {
       return '';
     }
 
-
     let replicaSet = description.setName;
     let serverTypePrompt = '';
     // TODO: replace with proper TopologyType constants - NODE-2973
@@ -348,7 +358,7 @@ export default class ShellInternalState {
         serverTypePrompt = '[primary]';
         break;
       case 'Sharded':
-        serverTypePrompt = '[mongos]';
+        serverTypePrompt = this.connectionInfo?.extraInfo?.is_atlas ? '[atlas proxy]' : '[mongos]';
         break;
       default:
         return '';