fix(shell-api): fix prompt for Atlas-proxy-style mongos MONGOSH-680 (#797)

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`.

Author: addaleax

packages/cli-repl/src/mongosh-repl.spec.ts

@@ -678,5 +678,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 > ');
+    });
   });
 });

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;
   }

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();

packages/shell-api/src/shell-internal-state.ts

@@ -236,6 +236,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);
   }
@@ -249,7 +260,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':
@@ -320,12 +331,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
@@ -342,7 +352,7 @@ export default class ShellInternalState {
         serverTypePrompt = '[primary]';
         break;
       case 'Sharded':
-        serverTypePrompt = '[mongos]';
+        serverTypePrompt = this.connectionInfo?.extraInfo?.is_atlas ? '[atlas proxy]' : '[mongos]';
         break;
       default:
         return '';