From 2d880715823dd94bde769d34cc030cc3dc9aa29b Mon Sep 17 00:00:00 2001 From: Massimiliano Marcon Date: Tue, 14 May 2024 13:53:08 +0200 Subject: [PATCH] Skunktendo 2024: automatically detect Local Atlas With this PR, the VS Code extension detects local atlas environments and shows them in the list of connections so that the user does not need to pass connection strings around. It does so by looking at the running docker containers, as that is how local atlas environments are provided. As added bonus, connections can now be open in Compass directy from VSCode if the protocol handler is registered. --- package-lock.json | 12 +++- package.json | 27 +++++++ src/commands/index.ts | 3 + src/connectionController.ts | 8 ++- src/mdbExtensionController.ts | 21 ++++++ src/storage/connectionStorage.ts | 118 ++++++++++++++++++++++++++++++- 6 files changed, 185 insertions(+), 4 deletions(-) diff --git a/package-lock.json b/package-lock.json index 35e3f803a..f6d834057 100644 --- a/package-lock.json +++ b/package-lock.json @@ -24,7 +24,7 @@ "bson-transpilers": "^2.2.0", "debug": "^4.3.4", "dotenv": "^16.3.1", - "express": "^4.19.2", + "jsonlines": "^0.1.1", "lodash": "^4.17.21", "micromatch": "^4.0.5", "mongodb": "^6.3.0", @@ -15658,6 +15658,11 @@ "graceful-fs": "^4.1.6" } }, + "node_modules/jsonlines": { + "version": "0.1.1", + "resolved": "https://registry.npmjs.org/jsonlines/-/jsonlines-0.1.1.tgz", + "integrity": "sha512-ekDrAGso79Cvf+dtm+mL8OBI2bmAOt3gssYs833De/C9NmIpWDWyUO4zPgB5x2/OhY366dkhgfPMYfwZF7yOZA==" + }, "node_modules/jsx-ast-utils": { "version": "3.3.4", "resolved": "https://registry.npmjs.org/jsx-ast-utils/-/jsx-ast-utils-3.3.4.tgz", @@ -35784,6 +35789,11 @@ "graceful-fs": "^4.1.6" } }, + "jsonlines": { + "version": "0.1.1", + "resolved": "https://registry.npmjs.org/jsonlines/-/jsonlines-0.1.1.tgz", + "integrity": "sha512-ekDrAGso79Cvf+dtm+mL8OBI2bmAOt3gssYs833De/C9NmIpWDWyUO4zPgB5x2/OhY366dkhgfPMYfwZF7yOZA==" + }, "jsx-ast-utils": { "version": "3.3.4", "resolved": "https://registry.npmjs.org/jsx-ast-utils/-/jsx-ast-utils-3.3.4.tgz", diff --git a/package.json b/package.json index ef3087a29..f2d9b4a21 100644 --- a/package.json +++ b/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", diff --git a/src/commands/index.ts b/src/commands/index.ts index ebef4dec1..7ef543f7c 100644 --- a/src/commands/index.ts +++ b/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; diff --git a/src/connectionController.ts b/src/connectionController.ts index 1572fd834..0eb5ba1e2 100644 --- a/src/connectionController.ts +++ b/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({ diff --git a/src/mdbExtensionController.ts b/src/mdbExtensionController.ts index 791d85287..a4fb49a5a 100644 --- a/src/mdbExtensionController.ts +++ b/src/mdbExtensionController.ts @@ -155,6 +155,13 @@ export default class MDBExtensionController implements vscode.Disposable { 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 @@ export default class MDBExtensionController implements vscode.Disposable { return true; } ); + this.registerCommand( + EXTENSION_COMMANDS.MDB_OPEN_IN_COMPASS, + async (element: ConnectionTreeItem): Promise => { + 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) => diff --git a/src/storage/connectionStorage.ts b/src/storage/connectionStorage.ts index 56bb2f166..67dbb23f7 100644 --- a/src/storage/connectionStorage.ts +++ b/src/storage/connectionStorage.ts @@ -15,6 +15,10 @@ import { 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 @@ type StoreConnectionInfoWithSecretStorageLocation = StoreConnectionInfo & 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 @@ export class ConnectionStorage { }) ); - 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 { + 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 { + 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) {