feat: add atlas-connect-cluster tool

Author: fmenezes

src/tools/atlas/metadata/connectCluster.ts

@@ -0,0 +1,89 @@
+import { z } from "zod";
+import { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
+import { AtlasToolBase } from "../atlasTool.js";
+import { ToolArgs, OperationType } from "../../tool.js";
+
+function generateSecurePassword(): string { // TODO: use a better password generator
+    return `pwdMcp${Math.floor(Math.random() * 100000)}`;
+}
+
+export class ConnectClusterTool extends AtlasToolBase {
+    protected name = "atlas-connect-cluster";
+    protected description = "Connect to MongoDB Atlas cluster";
+    protected operationType: OperationType = "metadata";
+    protected argsShape = {
+        projectId: z.string().describe("Atlas project ID"),
+        clusterName: z.string().describe("Atlas cluster name"),
+    };
+
+    protected async execute({ projectId, clusterName }: ToolArgs<typeof this.argsShape>): Promise<CallToolResult> {
+        const cluster = await this.session.apiClient.getCluster({
+            params: {
+                path: {
+                    groupId: projectId,
+                    clusterName,
+                },
+            },
+        });
+
+        if (!cluster) {
+            throw new Error("Cluster not found");
+        }
+
+        if (!cluster.connectionStrings?.standardSrv || !cluster.connectionStrings?.standard) {
+            throw new Error("Connection string not available");
+        }
+
+        const username = `usrMcp${Math.floor(Math.random() * 100000)}`;
+        const password = generateSecurePassword();
+
+        const expiryMs = 1000 * 60 * 60 * 12; // 12 hours
+        const expiryDate = new Date(Date.now() + expiryMs);
+
+        await this.session.apiClient.createDatabaseUser({
+            params: {
+                path: {
+                    groupId: projectId,
+                }
+            },
+            body: {
+                databaseName: "admin",
+                groupId: projectId,
+                roles: [
+                    {
+                        roleName: "readWriteAnyDatabase",
+                        databaseName: "admin",
+                    },
+                ],
+                scopes: [{type: "CLUSTER", name: clusterName}],
+                username,
+                password,
+                awsIAMType: "NONE",
+                ldapAuthType: "NONE",
+                oidcAuthType: "NONE",
+                x509Type: "NONE",
+                deleteAfterDate: expiryDate.toISOString(),
+            }
+        });
+
+        setTimeout(async () => { // disconnect after 12 hours
+            if (this.session.serviceProvider) {
+                await this.session.serviceProvider?.close(true);
+                this.session.serviceProvider = undefined;
+            }
+        }, expiryMs);
+
+        const connectionString = (cluster.connectionStrings.standardSrv || cluster.connectionStrings.standard || "").replace('://', `://${username}:${password}@`) + `?authSource=admin`;
+
+        await this.connectToMongoDB(connectionString);
+
+        return {
+            content: [
+                {
+                    type: "text",
+                    text: `Connected to cluster "${clusterName}"`,
+                },
+            ]
+        };
+    }
+}

src/tools/atlas/tools.ts

@@ -8,6 +8,7 @@ import { ListDBUsersTool } from "./read/listDBUsers.js";
 import { CreateDBUserTool } from "./create/createDBUser.js";
 import { CreateProjectTool } from "./create/createProject.js";
 import { ListOrganizationsTool } from "./read/listOrgs.js";
+import { ConnectClusterTool } from "./metadata/connectCluster.js";
 
 export const AtlasTools = [
     ListClustersTool,
@@ -20,4 +21,5 @@ export const AtlasTools = [
     CreateDBUserTool,
     CreateProjectTool,
     ListOrganizationsTool,
+    ConnectClusterTool
 ];

src/tools/mongodb/mongodbTool.ts

@@ -46,21 +46,4 @@ export abstract class MongoDBToolBase extends ToolBase {
 
         return super.handleError(error, args);
     }
-
-    protected async connectToMongoDB(connectionString: string): Promise<void> {
-        const provider = await NodeDriverServiceProvider.connect(connectionString, {
-            productDocsLink: "https://docs.mongodb.com/todo-mcp",
-            productName: "MongoDB MCP",
-            readConcern: {
-                level: this.config.connectOptions.readConcern,
-            },
-            readPreference: this.config.connectOptions.readPreference,
-            writeConcern: {
-                w: this.config.connectOptions.writeConcern,
-            },
-            timeoutMS: this.config.connectOptions.timeoutMS,
-        });
-
-        this.session.serviceProvider = provider;
-    }
 }

src/tools/tool.ts

@@ -7,6 +7,7 @@ import { mongoLogId } from "mongodb-log-writer";
 import { Telemetry } from "../telemetry/telemetry.js";
 import { type ToolEvent } from "../telemetry/types.js";
 import { UserConfig } from "../config.js";
+import { NodeDriverServiceProvider } from "@mongosh/service-provider-node-driver";
 
 export type ToolArgs<Args extends ZodRawShape> = z.objectOutputType<Args, ZodNever>;
 
@@ -123,4 +124,21 @@ export abstract class ToolBase {
             ],
         };
     }
+
+    protected async connectToMongoDB(connectionString: string): Promise<void> {
+        const provider = await NodeDriverServiceProvider.connect(connectionString, {
+            productDocsLink: "https://docs.mongodb.com/todo-mcp",
+            productName: "MongoDB MCP",
+            readConcern: {
+                level: this.config.connectOptions.readConcern,
+            },
+            readPreference: this.config.connectOptions.readPreference,
+            writeConcern: {
+                w: this.config.connectOptions.writeConcern,
+            },
+            timeoutMS: this.config.connectOptions.timeoutMS,
+        });
+
+        this.session.serviceProvider = provider;
+    }
 }

tests/integration/tools/atlas/clusters.test.ts

@@ -29,6 +29,24 @@ async function deleteAndWaitCluster(session: Session, projectId: string, cluster
     }
 }
 
+async function waitClusterState(session: Session, projectId: string, clusterName: string, state: string) {
+    while (true) {
+        const cluster = await session.apiClient.getCluster({
+            params: {
+                path: {
+                    groupId: projectId,
+                    clusterName: clusterName,
+                },
+            },
+        });
+        if (cluster?.stateName === state) {
+            return;
+        }
+        await sleep(1000);
+    }
+}
+
+
 describeWithAtlas("clusters", (integration) => {
     withProject(integration, ({ getProjectId }) => {
         const clusterName = "ClusterTest-" + randomId;
@@ -117,5 +135,45 @@ describeWithAtlas("clusters", (integration) => {
                 expect(response.content[1].text).toContain(`${clusterName} | `);
             });
         });
+
+        describe("atlas-connect-cluster", () => {
+            beforeAll(async () => {
+                const projectId = getProjectId();
+                await waitClusterState(integration.mcpServer().session, projectId, clusterName, "IDLE");
+                const cluster = await integration.mcpServer().session.apiClient.getCluster({
+                    params: {
+                        path: {
+                            groupId: projectId,
+                            clusterName: clusterName,
+                        },
+                    }
+                });
+
+                console.log("Cluster connection string: ", cluster?.connectionStrings?.standardSrv || cluster?.connectionStrings?.standard);
+            });
+
+            it("should have correct metadata", async () => {
+                const { tools } = await integration.mcpClient().listTools();
+                const connectCluster = tools.find((tool) => tool.name === "atlas-connect-cluster");
+
+                expectDefined(connectCluster);
+                expect(connectCluster.inputSchema.type).toBe("object");
+                expectDefined(connectCluster.inputSchema.properties);
+                expect(connectCluster.inputSchema.properties).toHaveProperty("projectId");
+                expect(connectCluster.inputSchema.properties).toHaveProperty("clusterName");
+            });
+
+            it("connects to cluster", async () => {
+                const projectId = getProjectId();
+
+                const response = (await integration.mcpClient().callTool({
+                    name: "atlas-connect-cluster",
+                    arguments: { projectId, clusterName },
+                })) as CallToolResult;
+                expect(response.content).toBeArray();
+                expect(response.content).toHaveLength(1);
+                expect(response.content[0].text).toContain(`Connected to cluster "${clusterName}"`);
+            });
+        });
     });
 });