diff --git a/src/logger.ts b/src/logger.ts index 354d8956..8c1e7ff2 100644 --- a/src/logger.ts +++ b/src/logger.ts @@ -17,6 +17,8 @@ export const LogId = { atlasDeleteDatabaseUserFailure: mongoLogId(1_001_002), atlasConnectFailure: mongoLogId(1_001_003), atlasInspectFailure: mongoLogId(1_001_004), + atlasConnectAttempt: mongoLogId(1_001_005), + atlasConnectSucceeded: mongoLogId(1_001_006), telemetryDisabled: mongoLogId(1_002_001), telemetryEmitFailure: mongoLogId(1_002_002), diff --git a/src/session.ts b/src/session.ts index 0b23883b..d6df810b 100644 --- a/src/session.ts +++ b/src/session.ts @@ -67,30 +67,27 @@ export class Session extends EventEmitter<{ } this.serviceProvider = undefined; } - if (!this.connectedAtlasCluster) { - this.emit("disconnect"); - return; - } - void this.apiClient - .deleteDatabaseUser({ - params: { - path: { - groupId: this.connectedAtlasCluster.projectId, - username: this.connectedAtlasCluster.username, - databaseName: "admin", + if (this.connectedAtlasCluster?.username && this.connectedAtlasCluster?.projectId) { + void 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)); - logger.error( - LogId.atlasDeleteDatabaseUserFailure, - "atlas-connect-cluster", - `Error deleting previous database user: ${error.message}` - ); - }); - this.connectedAtlasCluster = undefined; - + }) + .catch((err: unknown) => { + const error = err instanceof Error ? err : new Error(String(err)); + logger.error( + LogId.atlasDeleteDatabaseUserFailure, + "atlas-connect-cluster", + `Error deleting previous database user: ${error.message}` + ); + }); + this.connectedAtlasCluster = undefined; + } this.emit("disconnect"); } diff --git a/src/tools/atlas/metadata/connectCluster.ts b/src/tools/atlas/metadata/connectCluster.ts index 18970e24..a65913a6 100644 --- a/src/tools/atlas/metadata/connectCluster.ts +++ b/src/tools/atlas/metadata/connectCluster.ts @@ -11,6 +11,7 @@ const EXPIRY_MS = 1000 * 60 * 60 * 12; // 12 hours function sleep(ms: number): Promise { return new Promise((resolve) => setTimeout(resolve, ms)); } + export class ConnectClusterTool extends AtlasToolBase { protected name = "atlas-connect-cluster"; protected description = "Connect to MongoDB Atlas cluster"; @@ -20,9 +21,46 @@ export class ConnectClusterTool extends AtlasToolBase { clusterName: z.string().describe("Atlas cluster name"), }; - protected async execute({ projectId, clusterName }: ToolArgs): Promise { - await this.session.disconnect(); + private async queryConnection( + projectId: string, + clusterName: string + ): Promise<"connected" | "disconnected" | "connecting" | "connected-to-other-cluster" | "unknown"> { + if (!this.session.connectedAtlasCluster) { + if (this.session.serviceProvider) { + return "connected-to-other-cluster"; + } + return "disconnected"; + } + + if ( + this.session.connectedAtlasCluster.projectId !== projectId || + this.session.connectedAtlasCluster.clusterName !== clusterName + ) { + return "connected-to-other-cluster"; + } + + if (!this.session.serviceProvider) { + return "connecting"; + } + try { + await this.session.serviceProvider.runCommand("admin", { + ping: 1, + }); + + return "connected"; + } catch (err: unknown) { + const error = err instanceof Error ? err : new Error(String(err)); + logger.debug( + LogId.atlasConnectFailure, + "atlas-connect-cluster", + `error querying cluster: ${error.message}` + ); + return "unknown"; + } + } + + private async prepareClusterConnection(projectId: string, clusterName: string): Promise { const cluster = await inspectCluster(this.session.apiClient, projectId, clusterName); if (!cluster.connectionString) { @@ -81,14 +119,32 @@ export class ConnectClusterTool extends AtlasToolBase { cn.username = username; cn.password = password; cn.searchParams.set("authSource", "admin"); - const connectionString = cn.toString(); + return cn.toString(); + } + private async connectToCluster(projectId: string, clusterName: string, connectionString: string): Promise { let lastError: Error | undefined = undefined; - for (let i = 0; i < 20; i++) { + logger.debug( + LogId.atlasConnectAttempt, + "atlas-connect-cluster", + `attempting to connect to cluster: ${this.session.connectedAtlasCluster?.clusterName}` + ); + + // try to connect for about 5 minutes + for (let i = 0; i < 600; i++) { + if ( + !this.session.connectedAtlasCluster || + this.session.connectedAtlasCluster.projectId != projectId || + this.session.connectedAtlasCluster.clusterName != clusterName + ) { + throw new Error("Cluster connection aborted"); + } + try { - await this.session.connectToMongoDB(connectionString, this.config.connectOptions); lastError = undefined; + + await this.session.connectToMongoDB(connectionString, this.config.connectOptions); break; } catch (err: unknown) { const error = err instanceof Error ? err : new Error(String(err)); @@ -106,14 +162,94 @@ export class ConnectClusterTool extends AtlasToolBase { } if (lastError) { + if ( + this.session.connectedAtlasCluster?.projectId == projectId && + this.session.connectedAtlasCluster?.clusterName == clusterName && + this.session.connectedAtlasCluster?.username + ) { + void this.session.apiClient + .deleteDatabaseUser({ + params: { + path: { + groupId: this.session.connectedAtlasCluster.projectId, + username: this.session.connectedAtlasCluster.username, + databaseName: "admin", + }, + }, + }) + .catch((err: unknown) => { + const error = err instanceof Error ? err : new Error(String(err)); + logger.debug( + LogId.atlasConnectFailure, + "atlas-connect-cluster", + `error deleting database user: ${error.message}` + ); + }); + } + this.session.connectedAtlasCluster = undefined; throw lastError; } + logger.debug( + LogId.atlasConnectSucceeded, + "atlas-connect-cluster", + `connected to cluster: ${this.session.connectedAtlasCluster?.clusterName}` + ); + } + + protected async execute({ projectId, clusterName }: ToolArgs): Promise { + for (let i = 0; i < 60; i++) { + const state = await this.queryConnection(projectId, clusterName); + switch (state) { + case "connected": { + return { + content: [ + { + type: "text", + text: `Connected to cluster "${clusterName}".`, + }, + ], + }; + } + case "connecting": { + break; + } + case "connected-to-other-cluster": + case "disconnected": + case "unknown": + default: { + await this.session.disconnect(); + const connectionString = await this.prepareClusterConnection(projectId, clusterName); + + // try to connect for about 5 minutes asynchronously + void this.connectToCluster(projectId, clusterName, connectionString).catch((err: unknown) => { + const error = err instanceof Error ? err : new Error(String(err)); + logger.error( + LogId.atlasConnectFailure, + "atlas-connect-cluster", + `error connecting to cluster: ${error.message}` + ); + }); + break; + } + } + + await sleep(500); + } + return { content: [ { - type: "text", - text: `Connected to cluster "${clusterName}"`, + type: "text" as const, + text: `Attempting to connect to cluster "${clusterName}"...`, + }, + { + type: "text" as const, + text: `Warning: Provisioning a user and connecting to the cluster may take more time, please check again in a few seconds.`, + }, + { + type: "text" as const, + text: `Warning: Make sure your IP address was enabled in the allow list setting of the Atlas cluster.`, }, ], }; diff --git a/src/tools/mongodb/mongodbTool.ts b/src/tools/mongodb/mongodbTool.ts index f215f9a2..fe996a38 100644 --- a/src/tools/mongodb/mongodbTool.ts +++ b/src/tools/mongodb/mongodbTool.ts @@ -14,16 +14,25 @@ export abstract class MongoDBToolBase extends ToolBase { protected category: ToolCategory = "mongodb"; protected async ensureConnected(): Promise { - if (!this.session.serviceProvider && this.config.connectionString) { - try { - await this.connectToMongoDB(this.config.connectionString); - } catch (error) { - logger.error( - LogId.mongodbConnectFailure, - "mongodbTool", - `Failed to connect to MongoDB instance using the connection string from the config: ${error as string}` + if (!this.session.serviceProvider) { + if (this.session.connectedAtlasCluster) { + throw new MongoDBError( + ErrorCodes.NotConnectedToMongoDB, + `Attempting to connect to Atlas cluster "${this.session.connectedAtlasCluster.clusterName}", try again in a few seconds.` ); - throw new MongoDBError(ErrorCodes.MisconfiguredConnectionString, "Not connected to MongoDB."); + } + + if (this.config.connectionString) { + try { + await this.connectToMongoDB(this.config.connectionString); + } catch (error) { + logger.error( + LogId.mongodbConnectFailure, + "mongodbTool", + `Failed to connect to MongoDB instance using the connection string from the config: ${error as string}` + ); + throw new MongoDBError(ErrorCodes.MisconfiguredConnectionString, "Not connected to MongoDB."); + } } } diff --git a/tests/integration/tools/atlas/clusters.test.ts b/tests/integration/tools/atlas/clusters.test.ts index 8bb19bda..62bd422c 100644 --- a/tests/integration/tools/atlas/clusters.test.ts +++ b/tests/integration/tools/atlas/clusters.test.ts @@ -183,13 +183,27 @@ describeWithAtlas("clusters", (integration) => { 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}"`); + for (let i = 0; i < 10; i++) { + const response = (await integration.mcpClient().callTool({ + name: "atlas-connect-cluster", + arguments: { projectId, clusterName }, + })) as CallToolResult; + expect(response.content).toBeArray(); + expect(response.content.length).toBeGreaterThanOrEqual(1); + expect(response.content[0]?.type).toEqual("text"); + const c = response.content[0] as { text: string }; + if ( + c.text.includes("Cluster is already connected.") || + c.text.includes(`Connected to cluster "${clusterName}"`) + ) { + break; // success + } else { + expect(response.content[0]?.text).toContain( + `Attempting to connect to cluster "${clusterName}"...` + ); + } + await sleep(500); + } }); }); });