Add connection-id support to the auto-connect feature

Author: kraenhansen

packages/compass-e2e-tests/tests/auto-connect.test.ts

@@ -12,6 +12,7 @@ import * as Selectors from '../helpers/selectors';
 import os from 'os';
 import path from 'path';
 import { promises as fs } from 'fs';
+import { DEFAULT_CONNECTION_NAMES } from '../helpers/test-runner-context';
 
 const connectionStringSuccess = 'mongodb://127.0.0.1:27091/test';
 const connectionNameSuccess = 'Success';
@@ -21,6 +22,22 @@ const connectionNameUnreachable = 'Unreachable';
 const connectionStringInvalid = 'http://example.com';
 const connectionNameInvalid = 'Invalid';
 
+async function createDefaultConnectionAndClose(name: string | undefined) {
+  const compass = await init(name);
+  try {
+    const { browser } = compass;
+    await browser.setupDefaultConnections();
+    const connectionName = DEFAULT_CONNECTION_NAMES[0];
+    const connectionId = await browser.getConnectionIdByName(connectionName);
+    if (!connectionId) {
+      throw new Error('Expected a connection id');
+    }
+    return { connectionName, connectionId };
+  } finally {
+    await cleanup(compass);
+  }
+}
+
 describe('Automatically connecting from the command line', function () {
   let tmpdir: string;
   let i = 0;
@@ -127,6 +144,22 @@ describe('Automatically connecting from the command line', function () {
     }
   });
 
+  it('works with a connection id on the command line', async function () {
+    // Create the connection
+    const { connectionId, connectionName } =
+      await createDefaultConnectionAndClose(this.test?.fullTitle());
+    const compass = await init(this.test?.fullTitle(), {
+      wrapBinary: positionalArgs([connectionId]),
+      // reuse the same user directory so we'd get the same connections
+      firstRun: false,
+    });
+    try {
+      await waitForConnectionSuccessAndCheckConnection(compass, connectionName);
+    } finally {
+      await cleanup(compass);
+    }
+  });
+
   it('fails with an unreachable URL', async function () {
     const args = [
       `--file=${path.join(tmpdir, 'exported.json')}`,

packages/compass/src/main/auto-connect.ts

@@ -5,6 +5,7 @@ import { hasDisallowedConnectionStringOptions } from './validate-connection-stri
 import COMPASS_ICON from './icon';
 import { createLogger, mongoLogId } from '@mongodb-js/compass-logging';
 import { redactConnectionString } from 'mongodb-connection-string-url';
+import { UUID } from 'mongodb';
 const { log } = createLogger('COMPASS-AUTO-CONNECT-MAIN');
 
 export type AutoConnectPreferences = Partial<
@@ -22,7 +23,7 @@ export type AutoConnectPreferences = Partial<
 let autoConnectWindow: BrowserWindow['id'] | undefined = undefined;
 const browserWindowStates = new Map<
   BrowserWindow['id'],
-  { url?: string; disconnected?: boolean }
+  { url?: string; connectionId?: string; disconnected?: boolean }
 >();
 
 export function resetForTesting(): void {
@@ -73,6 +74,16 @@ export function registerMongoDbUrlForBrowserWindow(
   browserWindowStates.set(bw.id, { url });
 }
 
+export function registerConnectionIdForBrowserWindow(
+  bw: Pick<BrowserWindow, 'id'> | undefined | null,
+  connectionId: string
+): void {
+  if (!bw) {
+    return;
+  }
+  browserWindowStates.set(bw.id, { connectionId });
+}
+
 export const getConnectionStringFromArgs = (args?: string[]) => args?.[0];
 
 const shouldPreventAutoConnect = async ({
@@ -111,6 +122,14 @@ export async function getWindowAutoConnectPreferences(
 
   const windowState = browserWindowStates.get(id);
 
+  // Auto connect to any connection ID, right away
+  if (windowState?.connectionId) {
+    return Promise.resolve({
+      positionalArguments: [windowState.connectionId],
+      shouldAutoConnect: true,
+    });
+  }
+
   // Do not auto-connect in Compass windows other than the primary/first-created one.
   // Do auto-connect if this is a new window on macos that has been opened via open-url.
   if (autoConnectWindow !== id && !windowState?.url) {
@@ -148,6 +167,17 @@ export async function getWindowAutoConnectPreferences(
     trustedConnectionString,
   } = preferences.getPreferences();
 
+  // Auto connect to any connection ID, right away
+  if (
+    positionalArguments?.length === 1 &&
+    UUID.isValid(positionalArguments[0])
+  ) {
+    return Promise.resolve({
+      positionalArguments: [positionalArguments[0]],
+      shouldAutoConnect: true,
+    });
+  }
+
   // The about: accounts for webdriverio in the e2e tests appending the argument for every run
   if (
     !file &&

packages/connection-storage/src/compass-main-connection-storage.spec.ts

@@ -1379,6 +1379,19 @@ describe('ConnectionStorage', function () {
       );
     });
 
+    it('should return autoConnectInfo when an existing connection id is provided', async function () {
+      const connectionInfo = getConnectionInfo({ lastUsed: new Date() });
+      await writeFakeConnection(tmpDir, { connectionInfo });
+      const info = await connectionStorage.getAutoConnectInfo({
+        shouldAutoConnect: true,
+        positionalArguments: [connectionInfo.id],
+      });
+      expect(info?.id).to.equal(connectionInfo.id);
+      expect(info?.connectionOptions).to.deep.equal(
+        connectionInfo.connectionOptions
+      );
+    });
+
     context('when autoConnectInfo is available', function () {
       beforeEach(async function () {
         await connectionStorage.getAutoConnectInfo({

packages/connection-storage/src/compass-main-connection-storage.ts

@@ -298,6 +298,17 @@ class CompassMainConnectionStorage implements ConnectionStorage {
         password,
       });
     } else {
+      // Assuming the positional argument refers to an existing connection id
+      if (positionalArguments.length === 1) {
+        // Attempt to load it and return if found
+        const possibleConnection = await this.load({
+          id: positionalArguments[0],
+        });
+        if (possibleConnection) {
+          return possibleConnection;
+        }
+      }
+      // Determine if this is a valid connection string or a connection ID
       const connectionString = getConnectionStringFromArgs(positionalArguments);
       if (!connectionString) {
         throw new Error('Could not find a connection string');