feat: add support for private connection strings

Author: blva

src/common/atlas/cluster.ts

@@ -1,4 +1,8 @@
-import type { ClusterDescription20240805, FlexClusterDescription20241113 } from "./openapi.js";
+import type {
+    ClusterConnectionStrings,
+    ClusterDescription20240805,
+    FlexClusterDescription20241113,
+} from "./openapi.js";
 import type { ApiClient } from "./apiClient.js";
 import { LogId } from "../logger.js";
 import { ConnectionString } from "mongodb-connection-string-url";
@@ -19,6 +23,7 @@ export interface Cluster {
     state?: "IDLE" | "CREATING" | "UPDATING" | "DELETING" | "REPAIRING";
     mongoDBVersion?: string;
     connectionString?: string;
+    connectionStrings?: ClusterConnectionStrings;
     processIds?: Array<string>;
 }
 
@@ -31,6 +36,7 @@ export function formatFlexCluster(cluster: FlexClusterDescription20241113): Clus
         state: cluster.stateName,
         mongoDBVersion: cluster.mongoDBVersion,
         connectionString,
+        connectionStrings: cluster.connectionStrings,
         processIds: extractProcessIds(cluster.connectionStrings?.standard ?? ""),
     };
 }
@@ -74,10 +80,21 @@ export function formatCluster(cluster: ClusterDescription20240805): Cluster {
         state: cluster.stateName,
         mongoDBVersion: cluster.mongoDBVersion,
         connectionString,
+        connectionStrings: cluster.connectionStrings,
         processIds: extractProcessIds(cluster.connectionStrings?.standard ?? ""),
     };
 }
 
+export function getConnectionString(
+    connectionStrings: ClusterConnectionStrings,
+    connectionType: "standard" | "private"
+): string | undefined {
+    if (connectionType === "standard") {
+        return connectionStrings.standardSrv || connectionStrings.standard;
+    }
+    return connectionStrings.privateSrv || connectionStrings.private;
+}
+
 export async function inspectCluster(apiClient: ApiClient, projectId: string, clusterName: string): Promise<Cluster> {
     try {
         const cluster = await apiClient.getCluster({

src/tools/args.ts

@@ -68,6 +68,12 @@ export const AtlasArgs = {
 
     password: (): z.ZodString =>
         z.string().min(1, "Password is required").max(100, "Password must be 100 characters or less"),
+
+    connectionType: (): z.ZodDefault<z.ZodEnum<["standard", "private"]>> =>
+        z
+            .enum(["standard", "private"])
+            .default("standard")
+            .describe("Desired connection type (standard or private) to an Atlas cluster"),
 };
 
 function toEJSON<T extends object | undefined>(value: T): T {

src/tools/atlas/connect/connectCluster.ts

@@ -8,6 +8,7 @@ import { ensureCurrentIpInAccessList } from "../../../common/atlas/accessListUti
 import type { AtlasClusterConnectionInfo } from "../../../common/connectionManager.js";
 import { getDefaultRoleFromConfig } from "../../../common/atlas/roles.js";
 import { AtlasArgs } from "../../args.js";
+import { getConnectionString } from "../../../common/atlas/cluster.js";
 
 const addedIpAccessListMessage =
     "Note: Your current IP address has been added to the Atlas project's IP access list to enable secure connection.";
@@ -22,6 +23,7 @@ function sleep(ms: number): Promise<void> {
 export const ConnectClusterArgs = {
     projectId: AtlasArgs.projectId().describe("Atlas project ID"),
     clusterName: AtlasArgs.clusterName().describe("Atlas cluster name"),
+    connectionType: AtlasArgs.connectionType().optional(),
 };
 
 export class ConnectClusterTool extends AtlasToolBase {
@@ -69,12 +71,16 @@ export class ConnectClusterTool extends AtlasToolBase {
 
     private async prepareClusterConnection(
         projectId: string,
-        clusterName: string
+        clusterName: string,
+        connectionType: "standard" | "private"
     ): Promise<{ connectionString: string; atlas: AtlasClusterConnectionInfo }> {
         const cluster = await inspectCluster(this.session.apiClient, projectId, clusterName);
-
-        if (!cluster.connectionString) {
-            throw new Error("Connection string not available");
+        if (cluster.connectionStrings === undefined) {
+            throw new Error("Connection strings not available");
+        }
+        const connectionString = getConnectionString(cluster.connectionStrings, connectionType);
+        if (connectionString === undefined) {
+            throw new Error("Connection string not available for connection type: " + connectionType);
         }
 
         const username = `mcpUser${Math.floor(Math.random() * 100000)}`;
@@ -200,7 +206,11 @@ export class ConnectClusterTool extends AtlasToolBase {
         });
     }
 
-    protected async execute({ projectId, clusterName }: ToolArgs<typeof this.argsShape>): Promise<CallToolResult> {
+    protected async execute({
+        projectId,
+        clusterName,
+        connectionType,
+    }: ToolArgs<typeof this.argsShape>): Promise<CallToolResult> {
         const ipAccessListUpdated = await ensureCurrentIpInAccessList(this.session.apiClient, projectId);
         let createdUser = false;
 
@@ -239,7 +249,11 @@ export class ConnectClusterTool extends AtlasToolBase {
                 case "disconnected":
                 default: {
                     await this.session.disconnect();
-                    const { connectionString, atlas } = await this.prepareClusterConnection(projectId, clusterName);
+                    const { connectionString, atlas } = await this.prepareClusterConnection(
+                        projectId,
+                        clusterName,
+                        connectionType
+                    );
 
                     createdUser = true;
                     // try to connect for about 5 minutes asynchronously