feat: Atlas get performance advisor integration tests (#589)

Author: kylelai1

src/common/atlas/cluster.ts

@@ -1,25 +1,37 @@
 import type { ClusterDescription20240805, FlexClusterDescription20241113 } from "./openapi.js";
 import type { ApiClient } from "./apiClient.js";
 import { LogId } from "../logger.js";
+import { ConnectionString } from "mongodb-connection-string-url";
 
-const DEFAULT_PORT = "27017";
+type AtlasProcessId = `${string}:${number}`;
+
+function extractProcessIds(connectionString: string): Array<AtlasProcessId> {
+    if (!connectionString) {
+        return [];
+    }
+    const connectionStringUrl = new ConnectionString(connectionString);
+    return connectionStringUrl.hosts as Array<AtlasProcessId>;
+}
 export interface Cluster {
     name?: string;
     instanceType: "FREE" | "DEDICATED" | "FLEX";
     instanceSize?: string;
     state?: "IDLE" | "CREATING" | "UPDATING" | "DELETING" | "REPAIRING";
     mongoDBVersion?: string;
     connectionString?: string;
+    processIds?: Array<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 +65,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 +112,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<Array<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)}`
         );
     }
 }

src/common/atlas/performanceAdvisorUtils.ts

@@ -1,13 +1,15 @@
 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"];
 export type DropIndexSuggestion = components["schemas"]["DropIndexSuggestionsIndex"];
 export type SlowQueryLogMetrics = components["schemas"]["PerformanceAdvisorSlowQueryMetrics"];
 export type SlowQueryLog = components["schemas"]["PerformanceAdvisorSlowQuery"];
 
+export const DEFAULT_SLOW_QUERY_LOGS_LIMIT = 50;
+
 interface SuggestedIndexesResponse {
     content: components["schemas"]["PerformanceAdvisorResponse"];
 }
@@ -112,22 +114,35 @@ 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: [] };
+        }
+
+        const slowQueryPromises = processIds.map((processId) =>
+            apiClient.listSlowQueries({
+                params: {
+                    path: {
+                        groupId: projectId,
+                        processId,
+                    },
+                    query: {
+                        ...(since && { since: since.getTime() }),
+                        ...(namespaces && { namespaces: namespaces }),
+                        nLogs: DEFAULT_SLOW_QUERY_LOGS_LIMIT,
+                    },
                 },
-            },
-        });
+            })
+        );
+
+        const responses = await Promise.allSettled(slowQueryPromises);
+
+        const allSlowQueryLogs = responses.reduce((acc, response) => {
+            return acc.concat(response.status === "fulfilled" ? (response.value.slowQueries ?? []) : []);
+        }, [] as Array<SlowQueryLog>);
 
-        return { slowQueryLogs: response.slowQueries ?? [] };
+        return { slowQueryLogs: allSlowQueryLogs };
     } catch (err) {
         apiClient.logger.debug({
             id: LogId.atlasPaSlowQueryLogsFailure,

src/tools/atlas/atlasTool.ts

@@ -38,14 +38,29 @@ For more information on setting up API keys, visit: https://www.mongodb.com/docs
                 };
             }
 
+            if (statusCode === 402) {
+                return {
+                    content: [
+                        {
+                            type: "text",
+                            text: `Received a Payment Required API Error: ${error.message}
+
+Payment setup is required to perform this action in MongoDB Atlas.
+Please ensure that your payment method for your organization has been set up and is active.
+For more information on setting up payment, visit: https://www.mongodb.com/docs/atlas/billing/`,
+                        },
+                    ],
+                };
+            }
+
             if (statusCode === 403) {
                 return {
                     content: [
                         {
                             type: "text",
                             text: `Received a Forbidden API Error: ${error.message}
 
-You don't have sufficient permissions to perform this action in MongoDB Atlas
+You don't have sufficient permissions to perform this action in MongoDB Atlas.
 Please ensure your API key has the necessary roles assigned.
 For more information on Atlas API access roles, visit: https://www.mongodb.com/docs/atlas/api/service-accounts-overview/`,
                         },

src/tools/atlas/read/getPerformanceAdvisor.ts

@@ -8,6 +8,7 @@ import {
     getDropIndexSuggestions,
     getSchemaAdvice,
     getSlowQueries,
+    DEFAULT_SLOW_QUERY_LOGS_LIMIT,
 } from "../../../common/atlas/performanceAdvisorUtils.js";
 import { AtlasArgs } from "../../args.js";
 
@@ -20,8 +21,7 @@ const PerformanceAdvisorOperationType = z.enum([
 
 export class GetPerformanceAdvisorTool extends AtlasToolBase {
     public name = "atlas-get-performance-advisor";
-    protected description =
-        "Get MongoDB Atlas performance advisor recommendations, which includes the operations: suggested indexes, drop index suggestions, slow query logs, and schema suggestions";
+    protected description = `Get MongoDB Atlas performance advisor recommendations, which includes the operations: suggested indexes, drop index suggestions, schema suggestions, and a sample of the most recent (max ${DEFAULT_SLOW_QUERY_LOGS_LIMIT}) slow query logs`;
     public operationType: OperationType = "read";
     protected argsShape = {
         projectId: AtlasArgs.projectId().describe("Atlas project ID to get performance advisor recommendations"),
@@ -31,8 +31,11 @@ export class GetPerformanceAdvisorTool extends AtlasToolBase {
             .default(PerformanceAdvisorOperationType.options)
             .describe("Operations to get performance advisor recommendations"),
         since: z
-            .date()
-            .describe("Date to get slow query logs since. Only relevant for the slowQueryLogs operation.")
+            .string()
+            .datetime()
+            .describe(
+                "Date to get slow query logs since. Must be a string in ISO 8601 format. Only relevant for the slowQueryLogs operation."
+            )
             .optional(),
         namespaces: z
             .array(z.string())
@@ -49,34 +52,58 @@ export class GetPerformanceAdvisorTool extends AtlasToolBase {
     }: ToolArgs<typeof this.argsShape>): Promise<CallToolResult> {
         try {
             const [suggestedIndexesResult, dropIndexSuggestionsResult, slowQueryLogsResult, schemaSuggestionsResult] =
-                await Promise.all([
+                await Promise.allSettled([
                     operations.includes("suggestedIndexes")
                         ? getSuggestedIndexes(this.session.apiClient, projectId, clusterName)
                         : Promise.resolve(undefined),
                     operations.includes("dropIndexSuggestions")
                         ? getDropIndexSuggestions(this.session.apiClient, projectId, clusterName)
                         : Promise.resolve(undefined),
                     operations.includes("slowQueryLogs")
-                        ? getSlowQueries(this.session.apiClient, projectId, clusterName, since, namespaces)
+                        ? getSlowQueries(
+                              this.session.apiClient,
+                              projectId,
+                              clusterName,
+                              since ? new Date(since) : undefined,
+                              namespaces
+                          )
                         : Promise.resolve(undefined),
                     operations.includes("schemaSuggestions")
                         ? getSchemaAdvice(this.session.apiClient, projectId, clusterName)
                         : Promise.resolve(undefined),
                 ]);
 
+            const hasSuggestedIndexes =
+                suggestedIndexesResult.status === "fulfilled" &&
+                suggestedIndexesResult.value?.suggestedIndexes &&
+                suggestedIndexesResult.value.suggestedIndexes.length > 0;
+            const hasDropIndexSuggestions =
+                dropIndexSuggestionsResult.status === "fulfilled" &&
+                dropIndexSuggestionsResult.value?.hiddenIndexes &&
+                dropIndexSuggestionsResult.value?.redundantIndexes &&
+                dropIndexSuggestionsResult.value?.unusedIndexes &&
+                dropIndexSuggestionsResult.value.hiddenIndexes.length > 0 &&
+                dropIndexSuggestionsResult.value.redundantIndexes.length > 0 &&
+                dropIndexSuggestionsResult.value.unusedIndexes.length > 0;
+            const hasSlowQueryLogs =
+                slowQueryLogsResult.status === "fulfilled" &&
+                slowQueryLogsResult.value?.slowQueryLogs &&
+                slowQueryLogsResult.value.slowQueryLogs.length > 0;
+            const hasSchemaSuggestions =
+                schemaSuggestionsResult.status === "fulfilled" &&
+                schemaSuggestionsResult.value?.recommendations &&
+                schemaSuggestionsResult.value.recommendations.length > 0;
+
+            // Inserts the performance advisor data with the relevant section header if it exists
             const performanceAdvisorData = [
-                suggestedIndexesResult && suggestedIndexesResult?.suggestedIndexes?.length > 0
-                    ? `## Suggested Indexes\n${JSON.stringify(suggestedIndexesResult.suggestedIndexes)}`
-                    : "No suggested indexes found.",
-                dropIndexSuggestionsResult
-                    ? `## Drop Index Suggestions\n${JSON.stringify(dropIndexSuggestionsResult)}`
-                    : "No drop index suggestions found.",
-                slowQueryLogsResult && slowQueryLogsResult?.slowQueryLogs?.length > 0
-                    ? `## Slow Query Logs\n${JSON.stringify(slowQueryLogsResult.slowQueryLogs)}`
-                    : "No slow query logs found.",
-                schemaSuggestionsResult && schemaSuggestionsResult?.recommendations?.length > 0
-                    ? `## Schema Suggestions\n${JSON.stringify(schemaSuggestionsResult.recommendations)}`
-                    : "No schema suggestions found.",
+                `## Suggested Indexes\n${
+                    hasSuggestedIndexes
+                        ? `Note: The "Weight" field is measured in bytes, and represents the estimated number of bytes saved in disk reads per executed read query that would be saved by implementing an index suggestion. Please convert this to MB or GB for easier readability.\n${JSON.stringify(suggestedIndexesResult.value?.suggestedIndexes)}`
+                        : "No suggested indexes found."
+                }`,
+                `## Drop Index Suggestions\n${hasDropIndexSuggestions ? JSON.stringify(dropIndexSuggestionsResult.value) : "No drop index suggestions found."}`,
+                `## Slow Query Logs\n${hasSlowQueryLogs ? JSON.stringify(slowQueryLogsResult.value?.slowQueryLogs) : "No slow query logs found."}`,
+                `## Schema Suggestions\n${hasSchemaSuggestions ? JSON.stringify(schemaSuggestionsResult.value?.recommendations) : "No schema suggestions found."}`,
             ];
 
             if (performanceAdvisorData.length === 0) {

tests/integration/helpers.ts

@@ -41,6 +41,9 @@ export const defaultDriverOptions: DriverOptions = {
     ...driverOptions,
 };
 
+// Timeout in milliseconds for long running tests: defaults to 20 minutes
+export const DEFAULT_LONG_RUNNING_TEST_WAIT_TIMEOUT_MS = 1_200_000;
+
 export function setupIntegrationTest(
     getUserConfig: () => UserConfig,
     getDriverOptions: () => DriverOptions,

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 type { Session } from "../../../../src/common/session.js";
 
 export type IntegrationTestFunction = (integration: IntegrationTest) => void;
 
@@ -143,3 +144,83 @@ 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 assertClusterIsAvailable(
+    session: Session,
+    projectId: string,
+    clusterName: string
+): Promise<boolean> {
+    try {
+        await session.apiClient.getCluster({
+            params: {
+                path: {
+                    groupId: projectId,
+                    clusterName,
+                },
+            },
+        });
+        return true;
+    } catch {
+        return false;
+    }
+}
+
+export async function deleteAndWaitCluster(
+    session: Session,
+    projectId: string,
+    clusterName: string,
+    pollingInterval: number = 1000,
+    maxPollingIterations: number = 300
+): Promise<void> {
+    await session.apiClient.deleteCluster({
+        params: {
+            path: {
+                groupId: projectId,
+                clusterName,
+            },
+        },
+    });
+
+    for (let i = 0; i < maxPollingIterations; i++) {
+        const isAvailable = await assertClusterIsAvailable(session, projectId, clusterName);
+        if (!isAvailable) {
+            return;
+        }
+        await sleep(pollingInterval);
+    }
+    throw new Error(
+        `Cluster deletion timeout: ${clusterName} did not delete within ${maxPollingIterations} iterations`
+    );
+}
+
+export async function waitCluster(
+    session: Session,
+    projectId: string,
+    clusterName: string,
+    check: (cluster: ClusterDescription20240805) => boolean | Promise<boolean>,
+    pollingInterval: number = 1000,
+    maxPollingIterations: number = 300
+): Promise<void> {
+    for (let i = 0; i < maxPollingIterations; i++) {
+        const cluster = await session.apiClient.getCluster({
+            params: {
+                path: {
+                    groupId: projectId,
+                    clusterName,
+                },
+            },
+        });
+        if (await check(cluster)) {
+            return;
+        }
+        await sleep(pollingInterval);
+    }
+
+    throw new Error(
+        `Cluster wait timeout: ${clusterName} did not meet condition within ${maxPollingIterations} iterations`
+    );
+}