Retrieve process ID from standard connection string

kylelai1 · kylelai1 · commit 93d975786598 · 2025-09-23T23:04:38.000-04:00
diff --git a/src/common/atlas/cluster.ts b/src/common/atlas/cluster.ts
@@ -3,23 +3,45 @@ import type { ApiClient } from "./apiClient.js";
 import { LogId } from "../logger.js";
 
 const DEFAULT_PORT = "27017";
+
+function extractProcessIds(connectionString: string): Array<string> {
+    if (!connectionString || connectionString === "") return [];
+
+    // Extract host:port pairs from connection string
+    const matches = connectionString.match(/^mongodb:\/\/([^\/]+)/);
+    if (!matches) {
+        return [];
+    }
+
+    // matches[1] gives us the host:port pairs
+    const hostsString = matches[1];
+    const hosts = hostsString?.split(",") ?? [];
+
+    return hosts?.map((host) => {
+        const [hostname, port] = host.split(":");
+        return `${hostname}:${port || DEFAULT_PORT}`;
+    });
+}
 export interface Cluster {
     name?: string;
     instanceType: "FREE" | "DEDICATED" | "FLEX";
     instanceSize?: string;
     state?: "IDLE" | "CREATING" | "UPDATING" | "DELETING" | "REPAIRING";
     mongoDBVersion?: string;
     connectionString?: string;
+    processIds?: string[];
 }
 
 export function formatFlexCluster(cluster: FlexClusterDescription20241113): Cluster {
+    const connectionString = cluster.connectionStrings?.standardSrv || cluster.connectionStrings?.standard;
     return {
         name: cluster.name,
         instanceType: "FLEX",
         instanceSize: undefined,
         state: cluster.stateName,
         mongoDBVersion: cluster.mongoDBVersion,
-        connectionString: cluster.connectionStrings?.standardSrv || cluster.connectionStrings?.standard,
+        connectionString,
+        processIds: extractProcessIds(cluster.connectionStrings?.standard ?? ""),
     };
 }
 
@@ -53,14 +75,16 @@ export function formatCluster(cluster: ClusterDescription20240805): Cluster {
 
     const instanceSize = regionConfigs[0]?.instanceSize ?? "UNKNOWN";
     const clusterInstanceType = instanceSize === "M0" ? "FREE" : "DEDICATED";
+    const connectionString = cluster.connectionStrings?.standardSrv || cluster.connectionStrings?.standard;
 
     return {
         name: cluster.name,
         instanceType: clusterInstanceType,
         instanceSize: clusterInstanceType === "DEDICATED" ? instanceSize : undefined,
         state: cluster.stateName,
         mongoDBVersion: cluster.mongoDBVersion,
-        connectionString: cluster.connectionStrings?.standardSrv || cluster.connectionStrings?.standard,
+        connectionString,
+        processIds: extractProcessIds(cluster.connectionStrings?.standard ?? ""),
     };
 }
 
@@ -98,21 +122,17 @@ export async function inspectCluster(apiClient: ApiClient, projectId: string, cl
     }
 }
 
-export async function getProcessIdFromCluster(
+export async function getProcessIdsFromCluster(
     apiClient: ApiClient,
     projectId: string,
     clusterName: string
-): Promise<string> {
+): Promise<string[]> {
     try {
         const cluster = await inspectCluster(apiClient, projectId, clusterName);
-        if (!cluster.connectionString) {
-            throw new Error("No connection string available for cluster");
-        }
-        const url = new URL(cluster.connectionString);
-        return `${url.hostname}:${url.port || DEFAULT_PORT}`;
+        return cluster.processIds || [];
     } catch (error) {
         throw new Error(
-            `Failed to get processId from cluster: ${error instanceof Error ? error.message : String(error)}`
+            `Failed to get processIds from cluster: ${error instanceof Error ? error.message : String(error)}`
         );
     }
 }
diff --git a/src/common/atlas/performanceAdvisorUtils.ts b/src/common/atlas/performanceAdvisorUtils.ts
@@ -1,6 +1,6 @@
 import { LogId } from "../logger.js";
 import type { ApiClient } from "./apiClient.js";
-import { getProcessIdFromCluster } from "./cluster.js";
+import { getProcessIdsFromCluster } from "./cluster.js";
 import type { components } from "./openapi.js";
 
 export type SuggestedIndex = components["schemas"]["PerformanceAdvisorIndex"];
@@ -112,22 +112,36 @@ export async function getSlowQueries(
     namespaces?: Array<string>
 ): Promise<{ slowQueryLogs: Array<SlowQueryLog> }> {
     try {
-        const processId = await getProcessIdFromCluster(apiClient, projectId, clusterName);
+        const processIds = await getProcessIdsFromCluster(apiClient, projectId, clusterName);
 
-        const response = await apiClient.listSlowQueries({
-            params: {
-                path: {
-                    groupId: projectId,
-                    processId,
-                },
-                query: {
-                    ...(since && { since: since.getTime() }),
-                    ...(namespaces && { namespaces: namespaces }),
+        if (processIds.length === 0) {
+            return { slowQueryLogs: [] };
+        }
+
+        // Make parallel API calls for each process ID
+        const slowQueryPromises = processIds.map((processId) =>
+            apiClient.listSlowQueries({
+                params: {
+                    path: {
+                        groupId: projectId,
+                        processId,
+                    },
+                    query: {
+                        ...(since && { since: since.getTime() }),
+                        ...(namespaces && { namespaces: namespaces }),
+                    },
                 },
-            },
-        });
+            })
+        );
+
+        const responses = await Promise.all(slowQueryPromises);
+
+        // Combine all slow query logs from all process IDs
+        const allSlowQueryLogs = responses.reduce((acc, response) => {
+            return acc.concat(response.slowQueries ?? []);
+        }, [] as Array<SlowQueryLog>);
 
-        return { slowQueryLogs: response.slowQueries ?? [] };
+        return { slowQueryLogs: allSlowQueryLogs };
     } catch (err) {
         apiClient.logger.debug({
             id: LogId.atlasPaSlowQueryLogsFailure,
diff --git a/tests/integration/tools/atlas/atlasHelpers.ts b/tests/integration/tools/atlas/atlasHelpers.ts
@@ -1,10 +1,11 @@
 import { ObjectId } from "mongodb";
-import type { Group } from "../../../../src/common/atlas/openapi.js";
+import type { ClusterDescription20240805, Group } from "../../../../src/common/atlas/openapi.js";
 import type { ApiClient } from "../../../../src/common/atlas/apiClient.js";
 import type { IntegrationTest } from "../../helpers.js";
 import { setupIntegrationTest, defaultTestConfig, defaultDriverOptions } from "../../helpers.js";
 import type { SuiteCollector } from "vitest";
 import { afterAll, beforeAll, describe } from "vitest";
+import { Session } from "../../../../src/common/session.js";
 
 export type IntegrationTestFunction = (integration: IntegrationTest) => void;
 
@@ -143,3 +144,61 @@ async function createProject(apiClient: ApiClient): Promise<Group & Required<Pic
 
     return group as Group & Required<Pick<Group, "id">>;
 }
+
+export function sleep(ms: number): Promise<void> {
+    return new Promise((resolve) => setTimeout(resolve, ms));
+}
+
+export async function deleteAndWaitCluster(
+    session: Session,
+    projectId: string,
+    clusterName: string,
+    pollingInterval: number = 1000
+): Promise<void> {
+    await session.apiClient.deleteCluster({
+        params: {
+            path: {
+                groupId: projectId,
+                clusterName,
+            },
+        },
+    });
+    while (true) {
+        try {
+            await session.apiClient.getCluster({
+                params: {
+                    path: {
+                        groupId: projectId,
+                        clusterName,
+                    },
+                },
+            });
+            await sleep(pollingInterval);
+        } catch {
+            break;
+        }
+    }
+}
+
+export async function waitCluster(
+    session: Session,
+    projectId: string,
+    clusterName: string,
+    check: (cluster: ClusterDescription20240805) => boolean | Promise<boolean>,
+    pollingInterval: number = 1000
+): Promise<void> {
+    while (true) {
+        const cluster = await session.apiClient.getCluster({
+            params: {
+                path: {
+                    groupId: projectId,
+                    clusterName,
+                },
+            },
+        });
+        if (await check(cluster)) {
+            return;
+        }
+        await sleep(pollingInterval);
+    }
+}
diff --git a/tests/integration/tools/atlas/clusters.test.ts b/tests/integration/tools/atlas/clusters.test.ts
@@ -1,61 +1,16 @@
 import type { Session } from "../../../../src/common/session.js";
 import { expectDefined, getDataFromUntrustedContent, getResponseElements } from "../../helpers.js";
-import { describeWithAtlas, withProject, randomId, parseTable } from "./atlasHelpers.js";
-import type { ClusterDescription20240805 } from "../../../../src/common/atlas/openapi.js";
+import {
+    describeWithAtlas,
+    withProject,
+    randomId,
+    parseTable,
+    deleteAndWaitCluster,
+    waitCluster,
+    sleep,
+} from "./atlasHelpers.js";
 import { afterAll, beforeAll, describe, expect, it } from "vitest";
 
-function sleep(ms: number): Promise<void> {
-    return new Promise((resolve) => setTimeout(resolve, ms));
-}
-
-async function deleteAndWaitCluster(session: Session, projectId: string, clusterName: string): Promise<void> {
-    await session.apiClient.deleteCluster({
-        params: {
-            path: {
-                groupId: projectId,
-                clusterName,
-            },
-        },
-    });
-    while (true) {
-        try {
-            await session.apiClient.getCluster({
-                params: {
-                    path: {
-                        groupId: projectId,
-                        clusterName,
-                    },
-                },
-            });
-            await sleep(1000);
-        } catch {
-            break;
-        }
-    }
-}
-
-async function waitCluster(
-    session: Session,
-    projectId: string,
-    clusterName: string,
-    check: (cluster: ClusterDescription20240805) => boolean | Promise<boolean>
-): Promise<void> {
-    while (true) {
-        const cluster = await session.apiClient.getCluster({
-            params: {
-                path: {
-                    groupId: projectId,
-                    clusterName,
-                },
-            },
-        });
-        if (await check(cluster)) {
-            return;
-        }
-        await sleep(1000);
-    }
-}
-
 describeWithAtlas("clusters", (integration) => {
     withProject(integration, ({ getProjectId, getIpAddress }) => {
         const clusterName = "ClusterTest-" + randomId;
diff --git a/tests/integration/tools/atlas/performanceAdvisor.test.ts b/tests/integration/tools/atlas/performanceAdvisor.test.ts
