Skip to content

Commit e03b7a6

Browse files
committed
chore: move AtlasConnectionInfo to the connection manager
But keep deleting users outside, as it's out of scope of the connection manager as it's not related to MongoDB itself but to Atlas.
1 parent f49e2b0 commit e03b7a6

File tree

3 files changed

+44
-29
lines changed

3 files changed

+44
-29
lines changed

src/common/session.ts

Lines changed: 17 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@ import { Implementation } from "@modelcontextprotocol/sdk/types.js";
33
import logger, { LogId } from "./logger.js";
44
import EventEmitter from "events";
55
import {
6-
AnyConnectionState,
6+
AtlasClusterConnectionInfo,
77
ConnectionManager,
88
ConnectionSettings,
99
ConnectionStateConnected,
@@ -33,12 +33,6 @@ export class Session extends EventEmitter<SessionEvents> {
3333
name: string;
3434
version: string;
3535
};
36-
connectedAtlasCluster?: {
37-
username: string;
38-
projectId: string;
39-
clusterName: string;
40-
expiryDate: Date;
41-
};
4236

4337
constructor({ apiBaseUrl, apiClientId, apiClientSecret, connectionManager }: SessionOptions) {
4438
super();
@@ -70,20 +64,24 @@ export class Session extends EventEmitter<SessionEvents> {
7064
}
7165

7266
async disconnect(): Promise<void> {
67+
const currentConnection = this.connectionManager.currentConnectionState;
68+
const atlasCluster =
69+
currentConnection.tag === "connected" ? currentConnection.connectedAtlasCluster : undefined;
70+
7371
try {
7472
await this.connectionManager.disconnect();
7573
} catch (err: unknown) {
7674
const error = err instanceof Error ? err : new Error(String(err));
7775
logger.error(LogId.mongodbDisconnectFailure, "Error closing service provider:", error.message);
7876
}
7977

80-
if (this.connectedAtlasCluster?.username && this.connectedAtlasCluster?.projectId) {
78+
if (atlasCluster?.username && atlasCluster?.projectId) {
8179
void this.apiClient
8280
.deleteDatabaseUser({
8381
params: {
8482
path: {
85-
groupId: this.connectedAtlasCluster.projectId,
86-
username: this.connectedAtlasCluster.username,
83+
groupId: atlasCluster.projectId,
84+
username: atlasCluster.username,
8785
databaseName: "admin",
8886
},
8987
},
@@ -96,7 +94,6 @@ export class Session extends EventEmitter<SessionEvents> {
9694
`Error deleting previous database user: ${error.message}`
9795
);
9896
});
99-
this.connectedAtlasCluster = undefined;
10097
}
10198
}
10299

@@ -127,4 +124,13 @@ export class Session extends EventEmitter<SessionEvents> {
127124

128125
throw new MongoDBError(ErrorCodes.NotConnectedToMongoDB, "Not connected to MongoDB");
129126
}
127+
128+
get connectedAtlasCluster(): AtlasClusterConnectionInfo | undefined {
129+
const connectionState = this.connectionManager.currentConnectionState;
130+
if (connectionState.tag === "connected") {
131+
return connectionState.connectedAtlasCluster;
132+
}
133+
134+
return undefined;
135+
}
130136
}

src/tools/atlas/connect/connectCluster.ts

Lines changed: 27 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ import { generateSecurePassword } from "../../../helpers/generatePassword.js";
66
import logger, { LogId } from "../../../common/logger.js";
77
import { inspectCluster } from "../../../common/atlas/cluster.js";
88
import { ensureCurrentIpInAccessList } from "../../../common/atlas/accessListUtils.js";
9+
import { AtlasClusterConnectionInfo } from "../../../common/connectionManager.js";
910

1011
const EXPIRY_MS = 1000 * 60 * 60 * 12; // 12 hours
1112

@@ -27,7 +28,7 @@ export class ConnectClusterTool extends AtlasToolBase {
2728
clusterName: string
2829
): Promise<"connected" | "disconnected" | "connecting" | "connected-to-other-cluster" | "unknown"> {
2930
if (!this.session.connectedAtlasCluster) {
30-
if (this.session.isConnectedToMongoDB()) {
31+
if (this.session.isConnectedToMongoDB) {
3132
return "connected-to-other-cluster";
3233
}
3334
return "disconnected";
@@ -40,7 +41,7 @@ export class ConnectClusterTool extends AtlasToolBase {
4041
return "connected-to-other-cluster";
4142
}
4243

43-
if (!this.session.isConnectedToMongoDB()) {
44+
if (!this.session.isConnectedToMongoDB) {
4445
return "connecting";
4546
}
4647

@@ -61,7 +62,10 @@ export class ConnectClusterTool extends AtlasToolBase {
6162
}
6263
}
6364

64-
private async prepareClusterConnection(projectId: string, clusterName: string): Promise<string> {
65+
private async prepareClusterConnection(
66+
projectId: string,
67+
clusterName: string
68+
): Promise<[string, AtlasClusterConnectionInfo]> {
6569
const cluster = await inspectCluster(this.session.apiClient, projectId, clusterName);
6670

6771
if (!cluster.connectionString) {
@@ -109,7 +113,7 @@ export class ConnectClusterTool extends AtlasToolBase {
109113
},
110114
});
111115

112-
this.session.connectedAtlasCluster = {
116+
const connectedAtlasCluster = {
113117
username,
114118
projectId,
115119
clusterName,
@@ -120,10 +124,15 @@ export class ConnectClusterTool extends AtlasToolBase {
120124
cn.username = username;
121125
cn.password = password;
122126
cn.searchParams.set("authSource", "admin");
123-
return cn.toString();
127+
return [cn.toString(), connectedAtlasCluster];
124128
}
125129

126-
private async connectToCluster(projectId: string, clusterName: string, connectionString: string): Promise<void> {
130+
private async connectToCluster(
131+
projectId: string,
132+
clusterName: string,
133+
connectionString: string,
134+
atlas: AtlasClusterConnectionInfo
135+
): Promise<void> {
127136
let lastError: Error | undefined = undefined;
128137

129138
logger.debug(
@@ -145,7 +154,7 @@ export class ConnectClusterTool extends AtlasToolBase {
145154
try {
146155
lastError = undefined;
147156

148-
await this.session.connectToMongoDB({ connectionString, ...this.config.connectOptions });
157+
await this.session.connectToMongoDB({ connectionString, ...this.config.connectOptions, atlas });
149158
break;
150159
} catch (err: unknown) {
151160
const error = err instanceof Error ? err : new Error(String(err));
@@ -187,7 +196,6 @@ export class ConnectClusterTool extends AtlasToolBase {
187196
);
188197
});
189198
}
190-
this.session.connectedAtlasCluster = undefined;
191199
throw lastError;
192200
}
193201

@@ -221,17 +229,19 @@ export class ConnectClusterTool extends AtlasToolBase {
221229
case "disconnected":
222230
default: {
223231
await this.session.disconnect();
224-
const connectionString = await this.prepareClusterConnection(projectId, clusterName);
232+
const [connectionString, atlas] = await this.prepareClusterConnection(projectId, clusterName);
225233

226234
// try to connect for about 5 minutes asynchronously
227-
void this.connectToCluster(projectId, clusterName, connectionString).catch((err: unknown) => {
228-
const error = err instanceof Error ? err : new Error(String(err));
229-
logger.error(
230-
LogId.atlasConnectFailure,
231-
"atlas-connect-cluster",
232-
`error connecting to cluster: ${error.message}`
233-
);
234-
});
235+
void this.connectToCluster(projectId, clusterName, connectionString, atlas).catch(
236+
(err: unknown) => {
237+
const error = err instanceof Error ? err : new Error(String(err));
238+
logger.error(
239+
LogId.atlasConnectFailure,
240+
"atlas-connect-cluster",
241+
`error connecting to cluster: ${error.message}`
242+
);
243+
}
244+
);
235245
break;
236246
}
237247
}

src/tools/mongodb/mongodbTool.ts

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,6 @@ import { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
55
import { ErrorCodes, MongoDBError } from "../../common/errors.js";
66
import logger, { LogId } from "../../common/logger.js";
77
import { Server } from "../../server.js";
8-
import { AnyConnectionState } from "../../common/connectionManager.js";
98

109
export const DbOperationArgs = {
1110
database: z.string().describe("Database name"),

0 commit comments

Comments
 (0)