Skip to content

Commit 567d497

Browse files
cveticmcveticm
authored
chore: Support atlas connect via private and private endpoint connection strings (#673)
1 parent 790c569 commit 567d497

File tree

6 files changed

+62
-16
lines changed

6 files changed

+62
-16
lines changed

‎src/common/atlas/cluster.ts‎

Lines changed: 29 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,27 @@ export async function inspectCluster(apiClient: ApiClient, projectId: string, cl
112114
}
113115
}
114116

117+
/**
118+
* Returns a connection string for the specified connectionType.
119+
* For "privateEndpoint", it returns the first private endpoint connection string available.
120+
*/
121+
export function getConnectionString(
122+
connectionStrings: ClusterConnectionStrings,
123+
connectionType: "standard" | "private" | "privateEndpoint"
124+
): string | undefined {
125+
switch (connectionType) {
126+
case "standard":
127+
return connectionStrings.standardSrv || connectionStrings.standard;
128+
case "private":
129+
return connectionStrings.privateSrv || connectionStrings.private;
130+
case "privateEndpoint":
131+
return (
132+
connectionStrings.privateEndpoint?.[0]?.srvConnectionString ||
133+
connectionStrings.privateEndpoint?.[0]?.connectionString
134+
);
135+
}
136+
}
137+
115138
export async function getProcessIdsFromCluster(
116139
apiClient: ApiClient,
117140
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: 25 additions & 7 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().describe(
26+
"Type of connection (standard, private, or privateEndpoint) to an Atlas cluster"
27+
),
2528
};
2629

2730
export class ConnectClusterTool extends AtlasToolBase {
@@ -69,12 +72,19 @@ 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(
86+
`Connection string for connection type "${connectionType}" is not available. Please ensure this connection type is set up in Atlas. See https://www.mongodb.com/docs/atlas/connect-to-database-deployment/#connect-to-an-atlas-cluster.`
87+
);
7888
}
7989

8090
const username = `mcpUser${Math.floor(Math.random() * 100000)}`;
@@ -113,7 +123,7 @@ export class ConnectClusterTool extends AtlasToolBase {
113123
expiryDate,
114124
};
115125

116-
const cn = new URL(cluster.connectionString);
126+
const cn = new URL(connectionString);
117127
cn.username = username;
118128
cn.password = password;
119129
cn.searchParams.set("authSource", "admin");
@@ -200,7 +210,11 @@ export class ConnectClusterTool extends AtlasToolBase {
200210
});
201211
}
202212

203-
protected async execute({ projectId, clusterName }: ToolArgs<typeof this.argsShape>): Promise<CallToolResult> {
213+
protected async execute({
214+
projectId,
215+
clusterName,
216+
connectionType,
217+
}: ToolArgs<typeof this.argsShape>): Promise<CallToolResult> {
204218
const ipAccessListUpdated = await ensureCurrentIpInAccessList(this.session.apiClient, projectId);
205219
let createdUser = false;
206220

@@ -239,7 +253,11 @@ export class ConnectClusterTool extends AtlasToolBase {
239253
case "disconnected":
240254
default: {
241255
await this.session.disconnect();
242-
const { connectionString, atlas } = await this.prepareClusterConnection(projectId, clusterName);
256+
const { connectionString, atlas } = await this.prepareClusterConnection(
257+
projectId,
258+
clusterName,
259+
connectionType
260+
);
243261

244262
createdUser = true;
245263
// 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
),

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

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -150,16 +150,18 @@ describeWithAtlas("clusters", (integration) => {
150150
expectDefined(connectCluster.inputSchema.properties);
151151
expect(connectCluster.inputSchema.properties).toHaveProperty("projectId");
152152
expect(connectCluster.inputSchema.properties).toHaveProperty("clusterName");
153+
expect(connectCluster.inputSchema.properties).toHaveProperty("connectionType");
153154
});
154155

155156
it("connects to cluster", async () => {
156157
const projectId = getProjectId();
158+
const connectionType = "standard";
157159
let connected = false;
158160

159161
for (let i = 0; i < 10; i++) {
160162
const response = await integration.mcpClient().callTool({
161163
name: "atlas-connect-cluster",
162-
arguments: { projectId, clusterName },
164+
arguments: { projectId, clusterName, connectionType },
163165
});
164166

165167
const elements = getResponseElements(response.content);

0 commit comments

Comments
 (0)