Skip to content

Commit e3da69b

Browse files
cveticmcveticm
committed
MCP-259: support atlas connect via private and private endpoint connection strings
1 parent 790c569 commit e3da69b

File tree

5 files changed

+73
-16
lines changed

5 files changed

+73
-16
lines changed

‎src/common/atlas/cluster.ts‎

Lines changed: 31 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,8 @@
1-
import type { ClusterDescription20240805, FlexClusterDescription20241113 } from "./openapi.js";
1+
import type {
2+
ClusterConnectionStrings,
3+
ClusterDescription20240805,
4+
FlexClusterDescription20241113,
5+
} from "./openapi.js";
26
import type { ApiClient } from "./apiClient.js";
37
import { LogId } from "../logger.js";
48
import { ConnectionString } from "mongodb-connection-string-url";
@@ -18,19 +22,18 @@ export interface Cluster {
1822
instanceSize?: string;
1923
state?: "IDLE" | "CREATING" | "UPDATING" | "DELETING" | "REPAIRING";
2024
mongoDBVersion?: string;
21-
connectionString?: string;
25+
connectionStrings?: ClusterConnectionStrings;
2226
processIds?: Array<string>;
2327
}
2428

2529
export function formatFlexCluster(cluster: FlexClusterDescription20241113): Cluster {
26-
const connectionString = cluster.connectionStrings?.standardSrv || cluster.connectionStrings?.standard;
2730
return {
2831
name: cluster.name,
2932
instanceType: "FLEX",
3033
instanceSize: undefined,
3134
state: cluster.stateName,
3235
mongoDBVersion: cluster.mongoDBVersion,
33-
connectionString,
36+
connectionStrings: cluster.connectionStrings,
3437
processIds: extractProcessIds(cluster.connectionStrings?.standard ?? ""),
3538
};
3639
}
@@ -65,15 +68,14 @@ export function formatCluster(cluster: ClusterDescription20240805): Cluster {
6568

6669
const instanceSize = regionConfigs[0]?.instanceSize ?? "UNKNOWN";
6770
const clusterInstanceType = instanceSize === "M0" ? "FREE" : "DEDICATED";
68-
const connectionString = cluster.connectionStrings?.standardSrv || cluster.connectionStrings?.standard;
6971

7072
return {
7173
name: cluster.name,
7274
instanceType: clusterInstanceType,
7375
instanceSize: clusterInstanceType === "DEDICATED" ? instanceSize : undefined,
7476
state: cluster.stateName,
7577
mongoDBVersion: cluster.mongoDBVersion,
76-
connectionString,
78+
connectionStrings: cluster.connectionStrings,
7779
processIds: extractProcessIds(cluster.connectionStrings?.standard ?? ""),
7880
};
7981
}
@@ -112,6 +114,29 @@ export async function inspectCluster(apiClient: ApiClient, projectId: string, cl
112114
}
113115
}
114116

117+
// getConnectionString returns a connection string given a connectionType.
118+
// For "privateEndpoint", it returns the first private endpoint connection string available.
119+
export function getConnectionString(
120+
connectionStrings: ClusterConnectionStrings,
121+
connectionType: "standard" | "private" | "privateEndpoint"
122+
): string | undefined {
123+
if (connectionStrings === undefined) {
124+
return undefined;
125+
}
126+
127+
switch (connectionType) {
128+
case "standard":
129+
return connectionStrings.standardSrv || connectionStrings.standard;
130+
case "private":
131+
return connectionStrings.privateSrv || connectionStrings.private;
132+
case "privateEndpoint":
133+
return (
134+
connectionStrings.privateEndpoint?.[0]?.srvConnectionString ||
135+
connectionStrings.privateEndpoint?.[0]?.connectionString
136+
);
137+
}
138+
}
139+
115140
export async function getProcessIdsFromCluster(
116141
apiClient: ApiClient,
117142
projectId: string,

‎src/tools/args.ts‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,9 @@ export const AtlasArgs = {
4141
.max(64, "Cluster name must be 64 characters or less")
4242
.regex(ALLOWED_CLUSTER_NAME_CHARACTERS_REGEX, ALLOWED_CLUSTER_NAME_CHARACTERS_ERROR),
4343

44+
connectionType: (): z.ZodDefault<z.ZodEnum<["standard", "private", "privateEndpoint"]>> =>
45+
z.enum(["standard", "private", "privateEndpoint"]).default("standard"),
46+
4447
projectName: (): z.ZodString =>
4548
z
4649
.string()

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

Lines changed: 37 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@ import { type OperationType, type ToolArgs } from "../../tool.js";
33
import { AtlasToolBase } from "../atlasTool.js";
44
import { generateSecurePassword } from "../../../helpers/generatePassword.js";
55
import { LogId } from "../../../common/logger.js";
6-
import { inspectCluster } from "../../../common/atlas/cluster.js";
6+
import { getConnectionString, inspectCluster } from "../../../common/atlas/cluster.js";
77
import { ensureCurrentIpInAccessList } from "../../../common/atlas/accessListUtils.js";
88
import type { AtlasClusterConnectionInfo } from "../../../common/connectionManager.js";
99
import { getDefaultRoleFromConfig } from "../../../common/atlas/roles.js";
@@ -22,6 +22,9 @@ function sleep(ms: number): Promise<void> {
2222
export const ConnectClusterArgs = {
2323
projectId: AtlasArgs.projectId().describe("Atlas project ID"),
2424
clusterName: AtlasArgs.clusterName().describe("Atlas cluster name"),
25+
connectionType: AtlasArgs.connectionType()
26+
.optional()
27+
.describe("Desired connection type (standard, private, or privateEndpoint) to an Atlas cluster"),
2528
};
2629

2730
export class ConnectClusterTool extends AtlasToolBase {
@@ -69,12 +72,17 @@ export class ConnectClusterTool extends AtlasToolBase {
6972

7073
private async prepareClusterConnection(
7174
projectId: string,
72-
clusterName: string
75+
clusterName: string,
76+
connectionType: "standard" | "private" | "privateEndpoint" | undefined = "standard"
7377
): Promise<{ connectionString: string; atlas: AtlasClusterConnectionInfo }> {
7478
const cluster = await inspectCluster(this.session.apiClient, projectId, clusterName);
7579

76-
if (!cluster.connectionString) {
77-
throw new Error("Connection string not available");
80+
if (cluster.connectionStrings === undefined) {
81+
throw new Error("Connection strings not available");
82+
}
83+
const connectionString = getConnectionString(cluster.connectionStrings, connectionType);
84+
if (connectionString === undefined) {
85+
throw new Error(`Connection string for type "${connectionType}" not available`);
7886
}
7987

8088
const username = `mcpUser${Math.floor(Math.random() * 100000)}`;
@@ -113,13 +121,26 @@ export class ConnectClusterTool extends AtlasToolBase {
113121
expiryDate,
114122
};
115123

116-
const cn = new URL(cluster.connectionString);
124+
this.session.logger.debug({
125+
id: LogId.atlasConnectFailure,
126+
context: "atlas-connect-cluster",
127+
message: `Connection string received: ${connectionString}`,
128+
});
129+
const cn = new URL(connectionString);
117130
cn.username = username;
118131
cn.password = password;
119-
cn.searchParams.set("authSource", "admin");
132+
if (connectionType !== "privateEndpoint") {
133+
cn.searchParams.set("authSource", "admin");
134+
}
120135

121136
this.session.keychain.register(username, "user");
122137
this.session.keychain.register(password, "password");
138+
const thing = cn.toString();
139+
this.session.logger.debug({
140+
id: LogId.atlasConnectFailure,
141+
context: "atlas-connect-cluster",
142+
message: `>>>>>> Connection string used: ${thing}`,
143+
});
123144

124145
return { connectionString: cn.toString(), atlas: connectedAtlasCluster };
125146
}
@@ -200,7 +221,11 @@ export class ConnectClusterTool extends AtlasToolBase {
200221
});
201222
}
202223

203-
protected async execute({ projectId, clusterName }: ToolArgs<typeof this.argsShape>): Promise<CallToolResult> {
224+
protected async execute({
225+
projectId,
226+
clusterName,
227+
connectionType,
228+
}: ToolArgs<typeof this.argsShape>): Promise<CallToolResult> {
204229
const ipAccessListUpdated = await ensureCurrentIpInAccessList(this.session.apiClient, projectId);
205230
let createdUser = false;
206231

@@ -239,7 +264,11 @@ export class ConnectClusterTool extends AtlasToolBase {
239264
case "disconnected":
240265
default: {
241266
await this.session.disconnect();
242-
const { connectionString, atlas } = await this.prepareClusterConnection(projectId, clusterName);
267+
const { connectionString, atlas } = await this.prepareClusterConnection(
268+
projectId,
269+
clusterName,
270+
connectionType
271+
);
243272

244273
createdUser = true;
245274
// try to connect for about 5 minutes asynchronously

‎src/tools/atlas/read/inspectCluster.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,7 @@ export class InspectClusterTool extends AtlasToolBase {
3030
"Cluster details:",
3131
`Cluster Name | Cluster Type | Tier | State | MongoDB Version | Connection String
3232
----------------|----------------|----------------|----------------|----------------|----------------
33-
${formattedCluster.name || "Unknown"} | ${formattedCluster.instanceType} | ${formattedCluster.instanceSize || "N/A"} | ${formattedCluster.state || "UNKNOWN"} | ${formattedCluster.mongoDBVersion || "N/A"} | ${formattedCluster.connectionString || "N/A"}`
33+
${formattedCluster.name || "Unknown"} | ${formattedCluster.instanceType} | ${formattedCluster.instanceSize || "N/A"} | ${formattedCluster.state || "UNKNOWN"} | ${formattedCluster.mongoDBVersion || "N/A"} | ${formattedCluster.connectionStrings?.standardSrv || formattedCluster.connectionStrings?.standard || "N/A"}`
3434
),
3535
};
3636
}

‎src/tools/atlas/read/listClusters.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -105,7 +105,7 @@ ${rows}`,
105105
----------------|----------------|----------------|----------------|----------------|----------------
106106
${allClusters
107107
.map((formattedCluster) => {
108-
return `${formattedCluster.name || "Unknown"} | ${formattedCluster.instanceType} | ${formattedCluster.instanceSize || "N/A"} | ${formattedCluster.state || "UNKNOWN"} | ${formattedCluster.mongoDBVersion || "N/A"} | ${formattedCluster.connectionString || "N/A"}`;
108+
return `${formattedCluster.name || "Unknown"} | ${formattedCluster.instanceType} | ${formattedCluster.instanceSize || "N/A"} | ${formattedCluster.state || "UNKNOWN"} | ${formattedCluster.mongoDBVersion || "N/A"} | ${formattedCluster.connectionStrings?.standardSrv || formattedCluster.connectionStrings?.standard || "N/A"}`;
109109
})
110110
.join("\n")}`
111111
),

0 commit comments

Comments
 (0)