Skunktendo 2024: automatically detect Local Atlas

package.json

@@ -274,6 +274,10 @@
         "command": "mdb.copyConnectionString",
         "title": "Copy Connection String"
       },
+      {
+        "command": "mdb.openInCompass",
+        "title": "Open in Compass"
+      },
       {
         "command": "mdb.renameConnection",
         "title": "Rename Connection..."
@@ -429,6 +433,10 @@
       {
         "command": "mdb.dropStreamProcessor",
         "title": "Drop Stream Processor..."
+      },
+      {
+        "command": "mdb.reloadConnections",
+        "title": "MongoDB: Reload Connections"
       }
     ],
     "menus": {
@@ -503,6 +511,11 @@
           "when": "view == mongoDBConnectionExplorer && viewItem == connectedConnectionTreeItem",
           "group": "4@1"
         },
+        {
+          "command": "mdb.openInCompass",
+          "when": "view == mongoDBConnectionExplorer && viewItem == connectedConnectionTreeItem",
+          "group": "4@1"
+        },
         {
           "command": "mdb.disconnectFromConnectionTreeItem",
           "when": "view == mongoDBConnectionExplorer && viewItem == connectedConnectionTreeItem",
@@ -538,6 +551,11 @@
           "when": "view == mongoDBConnectionExplorer && viewItem == disconnectedConnectionTreeItem",
           "group": "3@1"
         },
+        {
+          "command": "mdb.openInCompass",
+          "when": "view == mongoDBConnectionExplorer && viewItem == disconnectedConnectionTreeItem",
+          "group": "3@1"
+        },
         {
           "command": "mdb.treeItemRemoveConnection",
           "when": "view == mongoDBConnectionExplorer && viewItem == disconnectedConnectionTreeItem",
@@ -772,6 +790,10 @@
           "command": "mdb.copyConnectionString",
           "when": "false"
         },
+        {
+          "command": "mdb.openInCompass",
+          "when": "false"
+        },
         {
           "command": "mdb.renameConnection",
           "when": "false"
@@ -899,6 +921,10 @@
         {
           "command": "mdb.dropStreamProcessor",
           "when": "false"
+        },
+        {
+          "command": "mdb.reloadConnections",
+          "when": "true"
         }
       ]
     },
@@ -1091,6 +1117,7 @@
     "bson-transpilers": "^2.2.0",
     "debug": "^4.3.4",
     "dotenv": "^16.3.1",
+    "jsonlines": "^0.1.1",
     "lodash": "^4.17.21",
     "micromatch": "^4.0.5",
     "mongodb": "^6.3.0",

src/commands/index.ts

@@ -73,6 +73,9 @@ enum EXTENSION_COMMANDS {
   MDB_START_STREAM_PROCESSOR = 'mdb.startStreamProcessor',
   MDB_STOP_STREAM_PROCESSOR = 'mdb.stopStreamProcessor',
   MDB_DROP_STREAM_PROCESSOR = 'mdb.dropStreamProcessor',
+
+  MDB_RELOAD_CONNECTIONS = 'mdb.reloadConnections',
+  MDB_OPEN_IN_COMPASS = 'mdb.openInCompass',
 }
 
 export default EXTENSION_COMMANDS;

src/connectionController.ts

@@ -163,10 +163,14 @@ export default class ConnectionController {
       this._connections[connection.id] = connection;
     }
 
-    if (loadedConnections.length) {
-      this.eventEmitter.emit(DataServiceEventTypes.CONNECTIONS_DID_CHANGE);
+    for (const connectionId of Object.keys(this._connections)) {
+      if (!loadedConnections.find(c => c.id === connectionId)) {
+        delete this._connections[connectionId];
+      }
     }
 
+    this.eventEmitter.emit(DataServiceEventTypes.CONNECTIONS_DID_CHANGE);
+
     // TODO: re-enable with fewer 'Saved Connections Loaded' events
     // https://jira.mongodb.org/browse/VSCODE-462
     /* this._telemetryService.trackSavedConnectionsLoaded({

src/mdbExtensionController.ts

@@ -155,6 +155,13 @@
   registerCommands = (): void => {
     // Register our extension's commands. These are the event handlers and
     // control the functionality of our extension.
+
+    // ------ CONNECTIONS ------ //
+    this.registerCommand(EXTENSION_COMMANDS.MDB_RELOAD_CONNECTIONS, () => {
+      this._connectionController.loadSavedConnections();
+      return Promise.resolve(true);
+    });
+
     // ------ CONNECTION ------ //
     this.registerCommand(EXTENSION_COMMANDS.MDB_OPEN_OVERVIEW_PAGE, () => {
       this._webviewController.openWebview(this._context);
@@ -350,6 +357,20 @@
         return true;
       }
     );
+    this.registerCommand(
+      EXTENSION_COMMANDS.MDB_OPEN_IN_COMPASS,
+      async (element: ConnectionTreeItem): Promise<boolean> => {
+        const connectionString =
+          this._connectionController.copyConnectionStringByConnectionId(
+            element.connectionId
+          );
+
+        await vscode.env.openExternal(vscode.Uri.parse(connectionString));
+        void vscode.window.showInformationMessage('Opening connection in Compass...');
+
+        return true;
+      }
+    );
     this.registerCommand(
       EXTENSION_COMMANDS.MDB_REMOVE_CONNECTION_TREE_VIEW,
       (element: ConnectionTreeItem) =>

src/storage/connectionStorage.ts

@@ -15,6 +15,10 @@
   StorageVariables,
 } from './storageController';
 
+import { spawn } from 'child_process';
+import jsonlines from 'jsonlines';
+import { URL } from 'url';
+
 const log = createLogger('connection storage');
 
 export interface StoreConnectionInfo {
@@ -34,6 +38,27 @@
 export type LoadedConnection = StoreConnectionInfoWithConnectionOptions &
   StoreConnectionInfoWithSecretStorageLocation;
 
+
+interface ContainerPort {
+  host: string;
+  port: Number;
+  containerPort: Number;
+  protocol: 'tcp' | 'udp';
+}
+
+interface Container {
+  id: string;
+  name: string;
+  ports: ContainerPort[];
+}
+
+interface _Container {
+  ID: string;
+  Names: string;
+  Image: string;
+  Ports: string;
+}
+
 export class ConnectionStorage {
   _storageController: StorageController;
 
@@ -202,7 +227,98 @@
       })
     );
 
-    return loadedConnections;
+    return loadedConnections.concat(await this.loadConnectionsToLocalInstances());
+  }
+
+  static parsePorts(portsString) {
+    // Given the docker port string, parse it and return an array of objects
+    // Example input:
+    // 0.0.0.0:27778->27017/tcp
+    // Example output:
+    // [{ host: '0.0.0.0', port: 27778, containerPort: 27017, protocol: 'tcp' }]
+    return portsString.split(",").map((portString) => {
+      const [host, container] = portString.split("->");
+      const [hostIp, hostPort] = host.split(":");
+      const [containerPort, protocol] = container.split("/");
+      return {
+        host: hostIp,
+        port: parseInt(hostPort),
+        containerPort: parseInt(containerPort),
+        protocol,
+      };
+    });
+  }
+
+  static toLoadedConnection(container: Container): LoadedConnection {
+    return {
+      id: container.id,
+      name: `LocalAtlas-${container.name}`,
+      storageLocation: StorageLocation.NONE,
+      secretStorageLocation: 'vscode.SecretStorage',
+      connectionOptions: {
+        connectionString: `mongodb://localhost:${container.ports[0].port}/?directConnection=true`
+      }
+    }
+  }
+
+  static async addCredentials(loadedConnection: LoadedConnection): Promise<LoadedConnection> {
+    return new Promise((resolve, reject) => {
+      const docker = spawn("docker", ["inspect", loadedConnection.id]);
+      let dockerInspectOutput = '';
+      docker.stdout.on('data', data => dockerInspectOutput += data);
+      docker.stdout.on('end', () => {
+        const parsedOutput = JSON.parse(dockerInspectOutput);
+        const env = parsedOutput.pop().Config.Env;
+
+        const credentials = env.reduce((acc, envVar) => {
+          const usernameMatch = envVar.match(/^MONGODB_INITDB_ROOT_USERNAME=(.*)$/);
+          const passwordMatch = envVar.match(/^MONGODB_INITDB_ROOT_PASSWORD=(.*)$/);
+          if (usernameMatch) {
+            acc.username = usernameMatch[1];
+          }
+          if (passwordMatch) {
+            acc.password = passwordMatch[1];
+          }
+          return acc;
+        }, {});
+
+        const { username, password } = credentials;
+
+        const connString = new URL(loadedConnection.connectionOptions.connectionString);
+        connString.username = username;
+        connString.password = password;
+        loadedConnection.connectionOptions.connectionString = connString.toString();
+        resolve(loadedConnection);
+      })
+    });
+  }
+
+  async loadConnectionsToLocalInstances(): Promise<LoadedConnection[]> {
+    const imageRegex = /^mongodb\/mongodb-atlas-local(:[a-zA-Z0-9\.-]+)?$/;
+    return new Promise((resolve, reject) => {
+      const docker = spawn("docker", ["ps", "--format", "json"]);
+      const parser = jsonlines.parse();
+      const containers: _Container[] = [];
+      docker.stdout.pipe(parser);
+      parser.on('data', function (data: _Container) {
+        containers.push(data);
+      });
+
+      parser.on('end', async function () {
+
+        let localInstances = containers
+          .filter((container: _Container) => imageRegex.test(container.Image))
+          .map((container) => ConnectionStorage.toLoadedConnection({
+            id: container.ID,
+            name: container.Names,
+            ports: ConnectionStorage.parsePorts(container.Ports),
+          }));
+
+        localInstances = await Promise.all(localInstances.map(li => ConnectionStorage.addCredentials(li)));
+
+        resolve(localInstances);
+      });
+    });
   }
 
   async removeConnection(connectionId: string) {