VSCODE-124: Connect with ssh to shell

package.json

@@ -608,18 +608,18 @@
     "@fortawesome/free-solid-svg-icons": "^5.13.0",
     "@fortawesome/react-fontawesome": "^0.1.9",
     "@leafygreen-ui/toggle": "^3.0.0",
-    "@mongosh/browser-runtime-electron": "0.0.1-alpha.15",
-    "@mongosh/service-provider-server": "0.0.1-alpha.15",
-    "@mongosh/shell-api": "0.0.1-alpha.15",
+    "@mongosh/browser-runtime-electron": "0.0.1-alpha.17",
+    "@mongosh/service-provider-server": "0.0.1-alpha.17",
+    "@mongosh/shell-api": "0.0.1-alpha.17",
     "analytics-node": "^3.4.0-beta.1",
     "bson": "^4.0.3",
     "classnames": "^2.2.6",
     "debug": "^4.1.1",
     "dotenv": "^8.2.0",
     "encoding": "^0.1.12",
     "mongodb-cloud-info": "^1.1.2",
-    "mongodb-connection-model": "^16.1.0",
-    "mongodb-data-service": "^16.7.0",
+    "mongodb-connection-model": "^16.1.2",
+    "mongodb-data-service": "^16.8.0",
     "mongodb-ns": "^2.2.0",
     "mongodb-schema": "^8.2.5",
     "react": "^16.13.1",

src/connectionController.ts

@@ -20,13 +20,13 @@ const MAX_CONNECTION_NAME_LENGTH = 512;
 export enum DataServiceEventTypes {
   CONNECTIONS_DID_CHANGE = 'CONNECTIONS_DID_CHANGE',
   ACTIVE_CONNECTION_CHANGED = 'ACTIVE_CONNECTION_CHANGED',
-  ACTIVE_CONNECTION_CHANGING = 'ACTIVE_CONNECTION_CHANGING',
+  ACTIVE_CONNECTION_CHANGING = 'ACTIVE_CONNECTION_CHANGING'
 }
 
 export enum ConnectionTypes {
   CONNECTION_FORM = 'CONNECTION_FORM',
   CONNECTION_STRING = 'CONNECTION_STRING',
-  CONNECTION_ID = 'CONNECTION_ID',
+  CONNECTION_ID = 'CONNECTION_ID'
 }
 
 export type SavedConnectionInformation = {
@@ -258,17 +258,25 @@ export default class ConnectionController {
     connectionModel: ConnectionModelType,
     connectionType: ConnectionTypes
   ): Promise<boolean> => {
-    const { driverUrl, instanceId } = connectionModel.getAttributes({
+    const {
+      driverUrl,
+      instanceId,
+      sshTunnelOptions
+    } = connectionModel.getAttributes({
       derived: true
     });
     const connectionId = uuidv4();
     const connectionInformation: SavedConnectionInformation = {
       connectionModel,
       driverUrl
     };
+    const name =
+      sshTunnelOptions.host && sshTunnelOptions.port
+        ? `${sshTunnelOptions.host}:${sshTunnelOptions.port}`
+        : instanceId;
     const savedConnection: SavedConnection = {
       id: connectionId,
-      name: instanceId,
+      name,
       // To begin we just store it on the session, the storage controller
       // handles changing this based on user preference.
       storageLocation: StorageScope.NONE
@@ -578,13 +586,13 @@ export default class ConnectionController {
     const connectionNameToRemove:
       | string
       | undefined = await vscode.window.showQuickPick(
-        connectionIds.map(
-          (id, index) => `${index + 1}: ${this._connections[id].name}`
-        ),
-        {
-          placeHolder: 'Choose a connection to remove...'
-        }
-      );
+      connectionIds.map(
+        (id, index) => `${index + 1}: ${this._connections[id].name}`
+      ),
+      {
+        placeHolder: 'Choose a connection to remove...'
+      }
+    );
 
     if (!connectionNameToRemove) {
       return Promise.resolve(false);

src/connectionModelType.ts

@@ -3,6 +3,7 @@ type ConnectionAttributes = {
   driverUrlWithSsh: string;
   driverOptions: any;
   instanceId: string;
+  sshTunnelOptions: any;
 };
 
 export type ConnectionModelType = {

src/mdbExtensionController.ts

@@ -399,23 +399,31 @@ export default class MDBExtensionController implements vscode.Disposable {
 
   public openMongoDBShell(): Promise<boolean> {
     let mdbConnectionString;
+
     if (this._connectionController) {
-      const activeConnectionDriverUrl = this._connectionController.getActiveConnectionDriverUrl();
-      mdbConnectionString = activeConnectionDriverUrl
-        ? activeConnectionDriverUrl
+      const activeConnectionModel = this._connectionController
+        .getActiveConnectionModel()
+        ?.getAttributes({ derived: true });
+
+      mdbConnectionString = activeConnectionModel
+        ? activeConnectionModel.driverUrlWithSsh
         : '';
     }
+
     if (!mdbConnectionString) {
       vscode.window.showErrorMessage(
         'You need to be connected before launching the MongoDB Shell.'
       );
+
       return Promise.resolve(false);
     }
+
     const mongoDBShell = vscode.window.createTerminal({
       name: 'MongoDB Shell',
       env: { MDB_CONNECTION_STRING: mdbConnectionString }
     });
     const shellCommand = vscode.workspace.getConfiguration('mdb').get('shell');
+
     mongoDBShell.sendText(
       `${shellCommand} $MDB_CONNECTION_STRING; unset MDB_CONNECTION_STRING`
     );

src/test/suite/extension.test.ts

@@ -1,45 +1,37 @@
 import * as assert from 'assert';
 import * as vscode from 'vscode';
-import { before, after } from 'mocha';
+import { beforeEach, afterEach } from 'mocha';
 import * as sinon from 'sinon';
-
+import Connection = require('mongodb-connection-model/lib/model');
 import MDBExtensionController from '../../mdbExtensionController';
-
 import { TestExtensionContext } from './stubs';
-import { TEST_DATABASE_URI } from './dbTestHelper';
-import ConnectionController from '../../connectionController';
-import { StorageController } from '../../storage';
-import { StatusView } from '../../views';
-import TelemetryController from '../../telemetry/telemetryController';
 
 suite('Extension Test Suite', () => {
   vscode.window.showInformationMessage('Starting tests...');
 
-  const disposables: vscode.Disposable[] = [];
   const mockExtensionContext = new TestExtensionContext();
   const mockMDBExtension = new MDBExtensionController(mockExtensionContext);
-  const mockStorageController = new StorageController(mockExtensionContext);
-  const testTelemetryController = new TelemetryController(
-    mockStorageController,
-    mockExtensionContext
-  );
-  const testConnectionController = new ConnectionController(
-    new StatusView(mockExtensionContext),
-    mockStorageController,
-    testTelemetryController
-  );
   const sandbox = sinon.createSandbox();
+
   let fakeShowErrorMessage: any;
+  let fakeGetActiveConnectionModel: any;
+  let createTerminalSpy: any;
 
-  before(() => {
+  beforeEach(() => {
     sandbox.stub(vscode.window, 'showInformationMessage');
+
     fakeShowErrorMessage = sandbox.stub(vscode.window, 'showErrorMessage');
+    fakeGetActiveConnectionModel = sandbox.stub(
+      mockMDBExtension._connectionController,
+      'getActiveConnectionModel'
+    );
+
+    createTerminalSpy = sinon.spy(vscode.window, 'createTerminal');
   });
 
-  after(() => {
-    disposables.forEach((d) => d.dispose());
-    disposables.length = 0;
+  afterEach(() => {
     sandbox.restore();
+    sinon.restore();
   });
 
   test('commands are registered in vscode', async () => {
@@ -102,54 +94,69 @@ suite('Extension Test Suite', () => {
   });
 
   test('openMongoDBShell should open a terminal with the active connection driver url', async () => {
-    try {
-      const succesfullyConnected = await testConnectionController.addNewConnectionStringAndConnect(
-        TEST_DATABASE_URI
-      );
+    const driverUri =
+      'mongodb://localhost:27018/?readPreference=primary&ssl=false';
 
-      assert(
-        succesfullyConnected === true,
-        'Expected a successful connection response.'
-      );
+    fakeGetActiveConnectionModel.returns(
+      new Connection({
+        hostname: 'localhost',
+        port: 27018
+      })
+    );
 
+    try {
       await mockMDBExtension.openMongoDBShell();
 
-      const spyActiveConnectionDriverUrl = sinon.spy(
-        testConnectionController,
-        'getActiveConnectionDriverUrl'
+      assert(createTerminalSpy.called);
+
+      const terminalOptions: vscode.TerminalOptions =
+        createTerminalSpy.firstCall.args[0];
+
+      assert(
+        terminalOptions.env?.MDB_CONNECTION_STRING === driverUri,
+        `Expected open terminal to set env var 'MDB_CONNECTION_STRING' to ${driverUri} found ${terminalOptions.env?.MDB_CONNECTION_STRING}`
       );
-      const createTerminalSpy = sinon.spy(vscode.window, 'createTerminal');
 
-      const checkResult = async () => {
-        await testConnectionController.disconnect();
+      await mockMDBExtension._connectionController.disconnect();
+      mockMDBExtension._connectionController.clearAllConnections();
+    } catch (e) {
+      assert(false);
+      return;
+    }
+  });
 
-        try {
-          assert(spyActiveConnectionDriverUrl.called);
-          assert(createTerminalSpy.called);
+  test('openMongoDBShell should open a terminal with ssh tunnel port injected', async () => {
+    const driverUri =
+      'mongodb://127.0.0.1:27017/?readPreference=primary&ssl=false';
 
-          const expectedUri =
-            'mongodb://localhost:27018/?readPreference=primary&ssl=false';
+    fakeGetActiveConnectionModel.returns(
+      new Connection({
+        hostname: '127.0.0.1',
+        sshTunnel: 'USER_PASSWORD',
+        sshTunnelHostname: 'my.ssh-server.com',
+        sshTunnelUsername: 'my-user',
+        sshTunnelPassword: 'password'
+      })
+    );
 
-          const terminalOptions: vscode.TerminalOptions = createTerminalSpy.firstCall.args[0];
+    try {
+      await mockMDBExtension.openMongoDBShell();
 
-          assert(
-            terminalOptions.env?.MDB_CONNECTION_STRING ===
-              expectedUri,
-            `Expected open terminal to set env var 'MDB_CONNECTION_STRING' to ${expectedUri} found ${
-              terminalOptions.env?.MDB_CONNECTION_STRING
-            }`
-          );
+      assert(createTerminalSpy.called);
 
-          testConnectionController.clearAllConnections();
-        } catch (e) {
-          assert(false);
-          return;
-        }
-      };
+      const terminalOptions: vscode.TerminalOptions =
+        createTerminalSpy.firstCall.args[0];
 
-      disposables.push(vscode.window.onDidOpenTerminal(checkResult));
-    } catch (error) {
+      assert(
+        terminalOptions.env?.MDB_CONNECTION_STRING !== driverUri,
+        `Expected open terminal to set env var 'MDB_CONNECTION_STRING' to driver url with ssh tunnel port injected found ${terminalOptions.env?.MDB_CONNECTION_STRING}`
+      );
+
+      await mockMDBExtension._connectionController.disconnect();
+      mockMDBExtension._connectionController.clearAllConnections();
+    } catch (e) {
       assert(false);
+      return;
     }
   });
 });

src/test/suite/telemetry/telemetryController.test.ts

@@ -172,7 +172,8 @@ suite('Telemetry Controller Test Suite', () => {
     sinon.assert.called(mockTrackPlaygroundCodeExecuted);
   });
 
-  test('track playground loaded and saved events', async () => {
+  test('track playground loaded and saved events', async function () {
+    this.timeout(3000);
     await loadAndSavePlayground(getDocUri('test.mongodb'));
 
     sinon.assert.called(mockTrackPlaygroundLoadedMethod);