feat: add atlas-connect-cluster tool (#131)

fmenezes · web-flow · commit 8680e3aaddaa · 2025-04-28T17:12:44.000+01:00
diff --git a/README.md b/README.md
@@ -104,6 +104,7 @@ You may experiment asking `Can you connect to my mongodb instance?`.
 - `atlas-list-clusters` - Lists MongoDB Atlas clusters
 - `atlas-inspect-cluster` - Inspect a specific MongoDB Atlas cluster
 - `atlas-create-free-cluster` - Create a free MongoDB Atlas cluster
+- `atlas-connect-cluster` - Connects to MongoDB Atlas cluster
 - `atlas-inspect-access-list` - Inspect IP/CIDR ranges with access to MongoDB Atlas clusters
 - `atlas-create-access-list` - Configure IP/CIDR access list for MongoDB Atlas clusters
 - `atlas-list-db-users` - List MongoDB Atlas database users
diff --git a/src/logger.ts b/src/logger.ts
@@ -11,6 +11,7 @@ export const LogId = {
     serverInitialized: mongoLogId(1_000_002),
 
     atlasCheckCredentials: mongoLogId(1_001_001),
+    atlasDeleteDatabaseUserFailure: mongoLogId(1_001_002),
 
     telemetryDisabled: mongoLogId(1_002_001),
     telemetryEmitFailure: mongoLogId(1_002_002),
@@ -22,6 +23,7 @@ export const LogId = {
     toolDisabled: mongoLogId(1_003_003),
 
     mongodbConnectFailure: mongoLogId(1_004_001),
+    mongodbDisconnectFailure: mongoLogId(1_004_002),
 } as const;
 
 abstract class LoggerBase {
diff --git a/src/session.ts b/src/session.ts
@@ -1,6 +1,7 @@
 import { NodeDriverServiceProvider } from "@mongosh/service-provider-node-driver";
 import { ApiClient, ApiClientCredentials } from "./common/atlas/apiClient.js";
 import { Implementation } from "@modelcontextprotocol/sdk/types.js";
+import logger, { LogId } from "./logger.js";
 import EventEmitter from "events";
 import { ConnectOptions } from "./config.js";
 
@@ -12,6 +13,7 @@ export interface SessionOptions {
 
 export class Session extends EventEmitter<{
     close: [];
+    disconnect: [];
 }> {
     sessionId?: string;
     serviceProvider?: NodeDriverServiceProvider;
@@ -20,6 +22,12 @@ export class Session extends EventEmitter<{
         name: string;
         version: string;
     };
+    connectedAtlasCluster?: {
+        username: string;
+        projectId: string;
+        clusterName: string;
+        expiryDate: Date;
+    };
 
     constructor({ apiBaseUrl, apiClientId, apiClientSecret }: SessionOptions) {
         super();
@@ -47,17 +55,47 @@ export class Session extends EventEmitter<{
         }
     }
 
-    async close(): Promise<void> {
+    async disconnect(): Promise<void> {
         if (this.serviceProvider) {
             try {
                 await this.serviceProvider.close(true);
-            } catch (error) {
-                console.error("Error closing service provider:", error);
+            } catch (err: unknown) {
+                const error = err instanceof Error ? err : new Error(String(err));
+                logger.error(LogId.mongodbDisconnectFailure, "Error closing service provider:", error.message);
             }
             this.serviceProvider = undefined;
+        }
+        if (!this.connectedAtlasCluster) {
+            this.emit("disconnect");
+            return;
+        }
+        try {
+            await this.apiClient.deleteDatabaseUser({
+                params: {
+                    path: {
+                        groupId: this.connectedAtlasCluster.projectId,
+                        username: this.connectedAtlasCluster.username,
+                        databaseName: "admin",
+                    },
+                },
+            });
+        } catch (err: unknown) {
+            const error = err instanceof Error ? err : new Error(String(err));
 
-            this.emit("close");
+            logger.error(
+                LogId.atlasDeleteDatabaseUserFailure,
+                "atlas-connect-cluster",
+                `Error deleting previous database user: ${error.message}`
+            );
         }
+        this.connectedAtlasCluster = undefined;
+
+        this.emit("disconnect");
+    }
+
+    async close(): Promise<void> {
+        await this.disconnect();
+        this.emit("close");
     }
 
     async connectToMongoDB(connectionString: string, connectOptions: ConnectOptions): Promise<void> {
diff --git a/src/tools/atlas/create/createFreeCluster.ts b/src/tools/atlas/create/createFreeCluster.ts
@@ -47,7 +47,10 @@ export class CreateFreeClusterTool extends AtlasToolBase {
         });
 
         return {
-            content: [{ type: "text", text: `Cluster "${name}" has been created in region "${region}".` }],
+            content: [
+                { type: "text", text: `Cluster "${name}" has been created in region "${region}".` },
+                { type: "text", text: `Double check your access lists to enable your current IP.` },
+            ],
         };
     }
 }
diff --git a/src/tools/atlas/metadata/connectCluster.ts b/src/tools/atlas/metadata/connectCluster.ts
@@ -0,0 +1,114 @@
+import { z } from "zod";
+import { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
+import { AtlasToolBase } from "../atlasTool.js";
+import { ToolArgs, OperationType } from "../../tool.js";
+import { randomBytes } from "crypto";
+import { promisify } from "util";
+
+const EXPIRY_MS = 1000 * 60 * 60 * 12; // 12 hours
+
+const randomBytesAsync = promisify(randomBytes);
+
+async function generateSecurePassword(): Promise<string> {
+    const buf = await randomBytesAsync(16);
+    const pass = buf.toString("base64url");
+    return pass;
+}
+
+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> {
+        await this.session.disconnect();
+
+        const cluster = await this.session.apiClient.getCluster({
+            params: {
+                path: {
+                    groupId: projectId,
+                    clusterName,
+                },
+            },
+        });
+
+        if (!cluster) {
+            throw new Error("Cluster not found");
+        }
+
+        const baseConnectionString = cluster.connectionStrings?.standardSrv || cluster.connectionStrings?.standard;
+
+        if (!baseConnectionString) {
+            throw new Error("Connection string not available");
+        }
+
+        const username = `mcpUser${Math.floor(Math.random() * 100000)}`;
+        const password = await generateSecurePassword();
+
+        const expiryDate = new Date(Date.now() + EXPIRY_MS);
+
+        const readOnly =
+            this.config.readOnly ||
+            (this.config.disabledTools?.includes("create") &&
+                this.config.disabledTools?.includes("update") &&
+                this.config.disabledTools?.includes("delete") &&
+                !this.config.disabledTools?.includes("read") &&
+                !this.config.disabledTools?.includes("metadata"));
+
+        const roleName = readOnly ? "readAnyDatabase" : "readWriteAnyDatabase";
+
+        await this.session.apiClient.createDatabaseUser({
+            params: {
+                path: {
+                    groupId: projectId,
+                },
+            },
+            body: {
+                databaseName: "admin",
+                groupId: projectId,
+                roles: [
+                    {
+                        roleName,
+                        databaseName: "admin",
+                    },
+                ],
+                scopes: [{ type: "CLUSTER", name: clusterName }],
+                username,
+                password,
+                awsIAMType: "NONE",
+                ldapAuthType: "NONE",
+                oidcAuthType: "NONE",
+                x509Type: "NONE",
+                deleteAfterDate: expiryDate.toISOString(),
+            },
+        });
+
+        this.session.connectedAtlasCluster = {
+            username,
+            projectId,
+            clusterName,
+            expiryDate,
+        };
+
+        const cn = new URL(baseConnectionString);
+        cn.username = username;
+        cn.password = password;
+        cn.searchParams.set("authSource", "admin");
+        const connectionString = cn.toString();
+
+        await this.session.connectToMongoDB(connectionString, this.config.connectOptions);
+
+        return {
+            content: [
+                {
+                    type: "text",
+                    text: `Connected to cluster "${clusterName}"`,
+                },
+            ],
+        };
+    }
+}
diff --git a/src/tools/atlas/tools.ts b/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,
 ];
diff --git a/tests/integration/tools/atlas/atlasHelpers.ts b/tests/integration/tools/atlas/atlasHelpers.ts
@@ -6,10 +6,6 @@ import { config } from "../../../../src/config.js";
 
 export type IntegrationTestFunction = (integration: IntegrationTest) => void;
 
-export function sleep(ms: number) {
-    return new Promise((resolve) => setTimeout(resolve, ms));
-}
-
 export function describeWithAtlas(name: string, fn: IntegrationTestFunction) {
     const testDefinition = () => {
         const integration = setupIntegrationTest(() => ({
diff --git a/tests/integration/tools/atlas/clusters.test.ts b/tests/integration/tools/atlas/clusters.test.ts
@@ -1,14 +1,18 @@
 import { Session } from "../../../../src/session.js";
 import { expectDefined } from "../../helpers.js";
-import { describeWithAtlas, withProject, sleep, randomId } from "./atlasHelpers.js";
+import { describeWithAtlas, withProject, randomId } from "./atlasHelpers.js";
 import { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
 
+function sleep(ms: number) {
+    return new Promise((resolve) => setTimeout(resolve, ms));
+}
+
 async function deleteAndWaitCluster(session: Session, projectId: string, clusterName: string) {
     await session.apiClient.deleteCluster({
         params: {
             path: {
                 groupId: projectId,
-                clusterName: clusterName,
+                clusterName,
             },
         },
     });
@@ -18,7 +22,7 @@ async function deleteAndWaitCluster(session: Session, projectId: string, cluster
                 params: {
                     path: {
                         groupId: projectId,
-                        clusterName: clusterName,
+                        clusterName,
                     },
                 },
             });
@@ -29,6 +33,23 @@ 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,
+                },
+            },
+        });
+        if (cluster?.stateName === state) {
+            return;
+        }
+        await sleep(1000);
+    }
+}
+
 describeWithAtlas("clusters", (integration) => {
     withProject(integration, ({ getProjectId }) => {
         const clusterName = "ClusterTest-" + randomId;
@@ -66,7 +87,7 @@ describeWithAtlas("clusters", (integration) => {
                     },
                 })) as CallToolResult;
                 expect(response.content).toBeArray();
-                expect(response.content).toHaveLength(1);
+                expect(response.content).toHaveLength(2);
                 expect(response.content[0].text).toContain("has been created");
             });
         });
@@ -117,5 +138,48 @@ 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");
+                await integration.mcpServer().session.apiClient.createProjectIpAccessList({
+                    params: {
+                        path: {
+                            groupId: projectId,
+                        },
+                    },
+                    body: [
+                        {
+                            comment: "MCP test",
+                            cidrBlock: "0.0.0.0/0",
+                        },
+                    ],
+                });
+            });
+
+            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}"`);
+            });
+        });
     });
 });

Original file line number	Diff line number	Diff line change
`@@ -47,7 +47,10 @@ export class CreateFreeClusterTool extends AtlasToolBase {`
`47`	`47`	`});`
`48`	`48`
`49`	`49`	`return {`
`50`		- content: [{ type: "text", text: `Cluster "${name}" has been created in region "${region}".` }],
	`50`	`+ content: [`
	`51`	+ { type: "text", text: `Cluster "${name}" has been created in region "${region}".` },
	`52`	+ { type: "text", text: `Double check your access lists to enable your current IP.` },
	`53`	`+ ],`
`51`	`54`	`};`
`52`	`55`	`}`
`53`	`56`	`}`