Skip to content

Commit ae1d6d0

Browse files
blvablva
committed
reformat
1 parent 78d430a commit ae1d6d0

File tree

2 files changed

+130
-64
lines changed

2 files changed

+130
-64
lines changed

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

Lines changed: 128 additions & 62 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,12 @@ import { inspectCluster } from "../../../common/atlas/cluster.js";
88
import { ensureCurrentIpInAccessList } from "../../../common/atlas/accessListUtils.js";
99
import { AtlasClusterConnectionInfo } from "../../../common/connectionManager.js";
1010

11-
const EXPIRY_MS = 1000 * 60 * 60 * 12; // 12 hours
11+
// Connection configuration constants
12+
const USER_EXPIRY_MS = 1000 * 60 * 60 * 12; // 12 hours
13+
const CONNECTION_RETRY_ATTEMPTS = 600; // 5 minutes (600 * 500ms = 5 minutes)
14+
const CONNECTION_RETRY_DELAY_MS = 500; // 500ms between retries
15+
const CONNECTION_CHECK_ATTEMPTS = 60; // 30 seconds (60 * 500ms = 30 seconds)
16+
const CONNECTION_CHECK_DELAY_MS = 500; // 500ms between connection state checks
1217

1318
function sleep(ms: number): Promise<void> {
1419
return new Promise((resolve) => setTimeout(resolve, ms));
@@ -23,6 +28,41 @@ export class ConnectClusterTool extends AtlasToolBase {
2328
clusterName: z.string().describe("Atlas cluster name"),
2429
};
2530

31+
private determineReadOnlyRole(): boolean {
32+
if (this.config.readOnly) return true;
33+
34+
const disabledTools = this.config.disabledTools || [];
35+
const hasWriteAccess = !disabledTools.includes("create") &&
36+
!disabledTools.includes("update") &&
37+
!disabledTools.includes("delete");
38+
const hasReadAccess = !disabledTools.includes("read") &&
39+
!disabledTools.includes("metadata");
40+
41+
return !hasWriteAccess && hasReadAccess;
42+
}
43+
44+
private isConnectedToOtherCluster(projectId: string, clusterName: string): boolean {
45+
return this.session.isConnectedToMongoDB &&
46+
(!this.session.connectedAtlasCluster ||
47+
this.session.connectedAtlasCluster.projectId !== projectId ||
48+
this.session.connectedAtlasCluster.clusterName !== clusterName);
49+
}
50+
51+
private getConnectionState(): "connected" | "connecting" | "disconnected" | "errored" {
52+
const state = this.session.connectionManager.currentConnectionState;
53+
switch (state.tag) {
54+
case "connected": return "connected";
55+
case "connecting": return "connecting";
56+
case "disconnected": return "disconnected";
57+
case "errored": return "errored";
58+
}
59+
}
60+
61+
private getErrorReason(): string | undefined {
62+
const state = this.session.connectionManager.currentConnectionState;
63+
return state.tag === "errored" ? state.errorReason : undefined;
64+
}
65+
2666
private queryConnection(
2767
projectId: string,
2868
clusterName: string
@@ -34,52 +74,40 @@ export class ConnectClusterTool extends AtlasToolBase {
3474
return "disconnected";
3575
}
3676

37-
const currentConectionState = this.session.connectionManager.currentConnectionState;
38-
if (
39-
this.session.connectedAtlasCluster.projectId !== projectId ||
40-
this.session.connectedAtlasCluster.clusterName !== clusterName
41-
) {
77+
if (this.isConnectedToOtherCluster(projectId, clusterName)) {
4278
return "connected-to-other-cluster";
4379
}
4480

45-
switch (currentConectionState.tag) {
81+
const connectionState = this.getConnectionState();
82+
switch (connectionState) {
4683
case "connecting":
4784
case "disconnected": // we might still be calling Atlas APIs and not attempted yet to connect to MongoDB, but we are still "connecting"
4885
return "connecting";
4986
case "connected":
5087
return "connected";
5188
case "errored":
89+
const errorReason = this.getErrorReason();
5290
this.session.logger.debug({
5391
id: LogId.atlasConnectFailure,
5492
context: "atlas-connect-cluster",
55-
message: `error querying cluster: ${currentConectionState.errorReason}`,
93+
message: `error querying cluster: ${errorReason || "unknown error"}`,
5694
});
5795
return "unknown";
5896
}
5997
}
6098

61-
private async prepareClusterConnection(
99+
private async createDatabaseUser(
62100
projectId: string,
63-
clusterName: string
64-
): Promise<{ connectionString: string; atlas: AtlasClusterConnectionInfo }> {
65-
const cluster = await inspectCluster(this.session.apiClient, projectId, clusterName);
66-
67-
if (!cluster.connectionString) {
68-
throw new Error("Connection string not available");
69-
}
70-
101+
clusterName: string,
102+
readOnly: boolean
103+
): Promise<{
104+
username: string;
105+
password: string;
106+
expiryDate: Date;
107+
}> {
71108
const username = `mcpUser${Math.floor(Math.random() * 100000)}`;
72109
const password = await generateSecurePassword();
73-
74-
const expiryDate = new Date(Date.now() + EXPIRY_MS);
75-
76-
const readOnly =
77-
this.config.readOnly ||
78-
(this.config.disabledTools?.includes("create") &&
79-
this.config.disabledTools?.includes("update") &&
80-
this.config.disabledTools?.includes("delete") &&
81-
!this.config.disabledTools?.includes("read") &&
82-
!this.config.disabledTools?.includes("metadata"));
110+
const expiryDate = new Date(Date.now() + USER_EXPIRY_MS);
83111

84112
const roleName = readOnly ? "readAnyDatabase" : "readWriteAnyDatabase";
85113

@@ -109,19 +137,52 @@ export class ConnectClusterTool extends AtlasToolBase {
109137
},
110138
});
111139

140+
return { username, password, expiryDate };
141+
}
142+
143+
private buildConnectionString(
144+
clusterConnectionString: string,
145+
username: string,
146+
password: string
147+
): string {
148+
const cn = new URL(clusterConnectionString);
149+
cn.username = username;
150+
cn.password = password;
151+
cn.searchParams.set("authSource", "admin");
152+
return cn.toString();
153+
}
154+
155+
private async prepareClusterConnection(
156+
projectId: string,
157+
clusterName: string
158+
): Promise<{ connectionString: string; atlas: AtlasClusterConnectionInfo }> {
159+
const cluster = await inspectCluster(this.session.apiClient, projectId, clusterName);
160+
161+
if (!cluster.connectionString) {
162+
throw new Error("Connection string not available");
163+
}
164+
165+
const readOnly = this.determineReadOnlyRole();
166+
const { username, password, expiryDate } = await this.createDatabaseUser(
167+
projectId,
168+
clusterName,
169+
readOnly
170+
);
171+
112172
const connectedAtlasCluster = {
113173
username,
114174
projectId,
115175
clusterName,
116176
expiryDate,
117177
};
118178

119-
const cn = new URL(cluster.connectionString);
120-
cn.username = username;
121-
cn.password = password;
122-
cn.searchParams.set("authSource", "admin");
179+
const connectionString = this.buildConnectionString(
180+
cluster.connectionString,
181+
username,
182+
password
183+
);
123184

124-
return { connectionString: cn.toString(), atlas: connectedAtlasCluster };
185+
return { connectionString, atlas: connectedAtlasCluster };
125186
}
126187

127188
private async connectToCluster(connectionString: string, atlas: AtlasClusterConnectionInfo): Promise<void> {
@@ -135,7 +196,7 @@ export class ConnectClusterTool extends AtlasToolBase {
135196
});
136197

137198
// try to connect for about 5 minutes
138-
for (let i = 0; i < 600; i++) {
199+
for (let i = 0; i < CONNECTION_RETRY_ATTEMPTS; i++) {
139200
try {
140201
lastError = undefined;
141202

@@ -152,7 +213,7 @@ export class ConnectClusterTool extends AtlasToolBase {
152213
message: `error connecting to cluster: ${error.message}`,
153214
});
154215

155-
await sleep(500); // wait for 500ms before retrying
216+
await sleep(CONNECTION_RETRY_DELAY_MS); // wait for 500ms before retrying
156217
}
157218

158219
if (
@@ -165,30 +226,7 @@ export class ConnectClusterTool extends AtlasToolBase {
165226
}
166227

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

@@ -200,9 +238,35 @@ export class ConnectClusterTool extends AtlasToolBase {
200238
});
201239
}
202240

241+
private async cleanupDatabaseUserOnFailure(atlas: AtlasClusterConnectionInfo): Promise<void> {
242+
const currentCluster = this.session.connectedAtlasCluster;
243+
if (currentCluster?.projectId === atlas.projectId &&
244+
currentCluster?.clusterName === atlas.clusterName &&
245+
currentCluster?.username) {
246+
try {
247+
await this.session.apiClient.deleteDatabaseUser({
248+
params: {
249+
path: {
250+
groupId: currentCluster.projectId,
251+
username: currentCluster.username,
252+
databaseName: "admin",
253+
},
254+
},
255+
});
256+
} catch (err: unknown) {
257+
const error = err instanceof Error ? err : new Error(String(err));
258+
this.session.logger.debug({
259+
id: LogId.atlasConnectFailure,
260+
context: "atlas-connect-cluster",
261+
message: `error deleting database user: ${error.message}`,
262+
});
263+
}
264+
}
265+
}
266+
203267
protected async execute({ projectId, clusterName }: ToolArgs<typeof this.argsShape>): Promise<CallToolResult> {
204268
await ensureCurrentIpInAccessList(this.session.apiClient, projectId);
205-
for (let i = 0; i < 60; i++) {
269+
for (let i = 0; i < CONNECTION_CHECK_ATTEMPTS; i++) {
206270
const state = this.queryConnection(projectId, clusterName);
207271
switch (state) {
208272
case "connected": {
@@ -226,19 +290,21 @@ export class ConnectClusterTool extends AtlasToolBase {
226290
const { connectionString, atlas } = await this.prepareClusterConnection(projectId, clusterName);
227291

228292
// try to connect for about 5 minutes asynchronously
229-
void this.connectToCluster(connectionString, atlas).catch((err: unknown) => {
293+
try {
294+
await this.connectToCluster(connectionString, atlas);
295+
} catch (err: unknown) {
230296
const error = err instanceof Error ? err : new Error(String(err));
231297
this.session.logger.error({
232298
id: LogId.atlasConnectFailure,
233299
context: "atlas-connect-cluster",
234300
message: `error connecting to cluster: ${error.message}`,
235301
});
236-
});
302+
}
237303
break;
238304
}
239305
}
240306

241-
await sleep(500);
307+
await sleep(CONNECTION_CHECK_DELAY_MS);
242308
}
243309

244310
return {

‎tests/unit/helpers/deviceId.test.ts‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -35,10 +35,10 @@ describe("deviceId", () => {
3535

3636
it("should correctly report initialization status", () => {
3737
expect(DeviceIdService.isInitialized()).toBe(false);
38-
38+
3939
DeviceIdService.init(testLogger);
4040
expect(DeviceIdService.isInitialized()).toBe(true);
41-
41+
4242
DeviceIdService.getInstance().close();
4343
expect(DeviceIdService.isInitialized()).toBe(false);
4444
});

0 commit comments

Comments
 (0)