Skip to content

Commit 27c52b4

Browse files
fmenezesfmenezes
authored
fix: turn atlas-connect-cluster async (#343)
1 parent 5ac06d0 commit 27c52b4

File tree

5 files changed

+204
-46
lines changed

5 files changed

+204
-46
lines changed

‎src/logger.ts

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,8 @@ export const LogId = {
1717
atlasDeleteDatabaseUserFailure: mongoLogId(1_001_002),
1818
atlasConnectFailure: mongoLogId(1_001_003),
1919
atlasInspectFailure: mongoLogId(1_001_004),
20+
atlasConnectAttempt: mongoLogId(1_001_005),
21+
atlasConnectSucceeded: mongoLogId(1_001_006),
2022

2123
telemetryDisabled: mongoLogId(1_002_001),
2224
telemetryEmitFailure: mongoLogId(1_002_002),

‎src/session.ts

Lines changed: 20 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -67,30 +67,27 @@ export class Session extends EventEmitter<{
6767
}
6868
this.serviceProvider = undefined;
6969
}
70-
if (!this.connectedAtlasCluster) {
71-
this.emit("disconnect");
72-
return;
73-
}
74-
void this.apiClient
75-
.deleteDatabaseUser({
76-
params: {
77-
path: {
78-
groupId: this.connectedAtlasCluster.projectId,
79-
username: this.connectedAtlasCluster.username,
80-
databaseName: "admin",
70+
if (this.connectedAtlasCluster?.username && this.connectedAtlasCluster?.projectId) {
71+
void this.apiClient
72+
.deleteDatabaseUser({
73+
params: {
74+
path: {
75+
groupId: this.connectedAtlasCluster.projectId,
76+
username: this.connectedAtlasCluster.username,
77+
databaseName: "admin",
78+
},
8179
},
82-
},
83-
})
84-
.catch((err: unknown) => {
85-
const error = err instanceof Error ? err : new Error(String(err));
86-
logger.error(
87-
LogId.atlasDeleteDatabaseUserFailure,
88-
"atlas-connect-cluster",
89-
`Error deleting previous database user: ${error.message}`
90-
);
91-
});
92-
this.connectedAtlasCluster = undefined;
93-
80+
})
81+
.catch((err: unknown) => {
82+
const error = err instanceof Error ? err : new Error(String(err));
83+
logger.error(
84+
LogId.atlasDeleteDatabaseUserFailure,
85+
"atlas-connect-cluster",
86+
`Error deleting previous database user: ${error.message}`
87+
);
88+
});
89+
this.connectedAtlasCluster = undefined;
90+
}
9491
this.emit("disconnect");
9592
}
9693

‎src/tools/atlas/metadata/connectCluster.ts

Lines changed: 143 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ const EXPIRY_MS = 1000 * 60 * 60 * 12; // 12 hours
1111
function sleep(ms: number): Promise<void> {
1212
return new Promise((resolve) => setTimeout(resolve, ms));
1313
}
14+
1415
export class ConnectClusterTool extends AtlasToolBase {
1516
protected name = "atlas-connect-cluster";
1617
protected description = "Connect to MongoDB Atlas cluster";
@@ -20,9 +21,46 @@ export class ConnectClusterTool extends AtlasToolBase {
2021
clusterName: z.string().describe("Atlas cluster name"),
2122
};
2223

23-
protected async execute({ projectId, clusterName }: ToolArgs<typeof this.argsShape>): Promise<CallToolResult> {
24-
await this.session.disconnect();
24+
private async queryConnection(
25+
projectId: string,
26+
clusterName: string
27+
): Promise<"connected" | "disconnected" | "connecting" | "connected-to-other-cluster" | "unknown"> {
28+
if (!this.session.connectedAtlasCluster) {
29+
if (this.session.serviceProvider) {
30+
return "connected-to-other-cluster";
31+
}
32+
return "disconnected";
33+
}
34+
35+
if (
36+
this.session.connectedAtlasCluster.projectId !== projectId ||
37+
this.session.connectedAtlasCluster.clusterName !== clusterName
38+
) {
39+
return "connected-to-other-cluster";
40+
}
41+
42+
if (!this.session.serviceProvider) {
43+
return "connecting";
44+
}
2545

46+
try {
47+
await this.session.serviceProvider.runCommand("admin", {
48+
ping: 1,
49+
});
50+
51+
return "connected";
52+
} catch (err: unknown) {
53+
const error = err instanceof Error ? err : new Error(String(err));
54+
logger.debug(
55+
LogId.atlasConnectFailure,
56+
"atlas-connect-cluster",
57+
`error querying cluster: ${error.message}`
58+
);
59+
return "unknown";
60+
}
61+
}
62+
63+
private async prepareClusterConnection(projectId: string, clusterName: string): Promise<string> {
2664
const cluster = await inspectCluster(this.session.apiClient, projectId, clusterName);
2765

2866
if (!cluster.connectionString) {
@@ -81,14 +119,32 @@ export class ConnectClusterTool extends AtlasToolBase {
81119
cn.username = username;
82120
cn.password = password;
83121
cn.searchParams.set("authSource", "admin");
84-
const connectionString = cn.toString();
122+
return cn.toString();
123+
}
85124

125+
private async connectToCluster(projectId: string, clusterName: string, connectionString: string): Promise<void> {
86126
let lastError: Error | undefined = undefined;
87127

88-
for (let i = 0; i < 20; i++) {
128+
logger.debug(
129+
LogId.atlasConnectAttempt,
130+
"atlas-connect-cluster",
131+
`attempting to connect to cluster: ${this.session.connectedAtlasCluster?.clusterName}`
132+
);
133+
134+
// try to connect for about 5 minutes
135+
for (let i = 0; i < 600; i++) {
136+
if (
137+
!this.session.connectedAtlasCluster ||
138+
this.session.connectedAtlasCluster.projectId != projectId ||
139+
this.session.connectedAtlasCluster.clusterName != clusterName
140+
) {
141+
throw new Error("Cluster connection aborted");
142+
}
143+
89144
try {
90-
await this.session.connectToMongoDB(connectionString, this.config.connectOptions);
91145
lastError = undefined;
146+
147+
await this.session.connectToMongoDB(connectionString, this.config.connectOptions);
92148
break;
93149
} catch (err: unknown) {
94150
const error = err instanceof Error ? err : new Error(String(err));
@@ -106,14 +162,94 @@ export class ConnectClusterTool extends AtlasToolBase {
106162
}
107163

108164
if (lastError) {
165+
if (
166+
this.session.connectedAtlasCluster?.projectId == projectId &&
167+
this.session.connectedAtlasCluster?.clusterName == clusterName &&
168+
this.session.connectedAtlasCluster?.username
169+
) {
170+
void this.session.apiClient
171+
.deleteDatabaseUser({
172+
params: {
173+
path: {
174+
groupId: this.session.connectedAtlasCluster.projectId,
175+
username: this.session.connectedAtlasCluster.username,
176+
databaseName: "admin",
177+
},
178+
},
179+
})
180+
.catch((err: unknown) => {
181+
const error = err instanceof Error ? err : new Error(String(err));
182+
logger.debug(
183+
LogId.atlasConnectFailure,
184+
"atlas-connect-cluster",
185+
`error deleting database user: ${error.message}`
186+
);
187+
});
188+
}
189+
this.session.connectedAtlasCluster = undefined;
109190
throw lastError;
110191
}
111192

193+
logger.debug(
194+
LogId.atlasConnectSucceeded,
195+
"atlas-connect-cluster",
196+
`connected to cluster: ${this.session.connectedAtlasCluster?.clusterName}`
197+
);
198+
}
199+
200+
protected async execute({ projectId, clusterName }: ToolArgs<typeof this.argsShape>): Promise<CallToolResult> {
201+
for (let i = 0; i < 60; i++) {
202+
const state = await this.queryConnection(projectId, clusterName);
203+
switch (state) {
204+
case "connected": {
205+
return {
206+
content: [
207+
{
208+
type: "text",
209+
text: `Connected to cluster "${clusterName}".`,
210+
},
211+
],
212+
};
213+
}
214+
case "connecting": {
215+
break;
216+
}
217+
case "connected-to-other-cluster":
218+
case "disconnected":
219+
case "unknown":
220+
default: {
221+
await this.session.disconnect();
222+
const connectionString = await this.prepareClusterConnection(projectId, clusterName);
223+
224+
// try to connect for about 5 minutes asynchronously
225+
void this.connectToCluster(projectId, clusterName, connectionString).catch((err: unknown) => {
226+
const error = err instanceof Error ? err : new Error(String(err));
227+
logger.error(
228+
LogId.atlasConnectFailure,
229+
"atlas-connect-cluster",
230+
`error connecting to cluster: ${error.message}`
231+
);
232+
});
233+
break;
234+
}
235+
}
236+
237+
await sleep(500);
238+
}
239+
112240
return {
113241
content: [
114242
{
115-
type: "text",
116-
text: `Connected to cluster "${clusterName}"`,
243+
type: "text" as const,
244+
text: `Attempting to connect to cluster "${clusterName}"...`,
245+
},
246+
{
247+
type: "text" as const,
248+
text: `Warning: Provisioning a user and connecting to the cluster may take more time, please check again in a few seconds.`,
249+
},
250+
{
251+
type: "text" as const,
252+
text: `Warning: Make sure your IP address was enabled in the allow list setting of the Atlas cluster.`,
117253
},
118254
],
119255
};

‎src/tools/mongodb/mongodbTool.ts

Lines changed: 18 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -14,16 +14,25 @@ export abstract class MongoDBToolBase extends ToolBase {
1414
protected category: ToolCategory = "mongodb";
1515

1616
protected async ensureConnected(): Promise<NodeDriverServiceProvider> {
17-
if (!this.session.serviceProvider && this.config.connectionString) {
18-
try {
19-
await this.connectToMongoDB(this.config.connectionString);
20-
} catch (error) {
21-
logger.error(
22-
LogId.mongodbConnectFailure,
23-
"mongodbTool",
24-
`Failed to connect to MongoDB instance using the connection string from the config: ${error as string}`
17+
if (!this.session.serviceProvider) {
18+
if (this.session.connectedAtlasCluster) {
19+
throw new MongoDBError(
20+
ErrorCodes.NotConnectedToMongoDB,
21+
`Attempting to connect to Atlas cluster "${this.session.connectedAtlasCluster.clusterName}", try again in a few seconds.`
2522
);
26-
throw new MongoDBError(ErrorCodes.MisconfiguredConnectionString, "Not connected to MongoDB.");
23+
}
24+
25+
if (this.config.connectionString) {
26+
try {
27+
await this.connectToMongoDB(this.config.connectionString);
28+
} catch (error) {
29+
logger.error(
30+
LogId.mongodbConnectFailure,
31+
"mongodbTool",
32+
`Failed to connect to MongoDB instance using the connection string from the config: ${error as string}`
33+
);
34+
throw new MongoDBError(ErrorCodes.MisconfiguredConnectionString, "Not connected to MongoDB.");
35+
}
2736
}
2837
}
2938

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

Lines changed: 21 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -183,13 +183,27 @@ describeWithAtlas("clusters", (integration) => {
183183
it("connects to cluster", async () => {
184184
const projectId = getProjectId();
185185

186-
const response = (await integration.mcpClient().callTool({
187-
name: "atlas-connect-cluster",
188-
arguments: { projectId, clusterName },
189-
})) as CallToolResult;
190-
expect(response.content).toBeArray();
191-
expect(response.content).toHaveLength(1);
192-
expect(response.content[0]?.text).toContain(`Connected to cluster "${clusterName}"`);
186+
for (let i = 0; i < 10; i++) {
187+
const response = (await integration.mcpClient().callTool({
188+
name: "atlas-connect-cluster",
189+
arguments: { projectId, clusterName },
190+
})) as CallToolResult;
191+
expect(response.content).toBeArray();
192+
expect(response.content.length).toBeGreaterThanOrEqual(1);
193+
expect(response.content[0]?.type).toEqual("text");
194+
const c = response.content[0] as { text: string };
195+
if (
196+
c.text.includes("Cluster is already connected.") ||
197+
c.text.includes(`Connected to cluster "${clusterName}"`)
198+
) {
199+
break; // success
200+
} else {
201+
expect(response.content[0]?.text).toContain(
202+
`Attempting to connect to cluster "${clusterName}"...`
203+
);
204+
}
205+
await sleep(500);
206+
}
193207
});
194208
});
195209
});

0 commit comments

Comments
 (0)