feat(explain): Add support for custom verbosity to the explain plan tool (#573)

Author: kmruiz

src/tools/mongodb/metadata/explain.ts

@@ -4,7 +4,6 @@ import type { ToolArgs, OperationType } from "../../tool.js";
 import { formatUntrustedData } from "../../tool.js";
 import { z } from "zod";
 import type { Document } from "mongodb";
-import { ExplainVerbosity } from "mongodb";
 import { AggregateArgs } from "../read/aggregate.js";
 import { FindArgs } from "../read/find.js";
 import { CountArgs } from "../read/count.js";
@@ -34,16 +33,22 @@ export class ExplainTool extends MongoDBToolBase {
                 ])
             )
             .describe("The method and its arguments to run"),
+        verbosity: z
+            .enum(["queryPlanner", "queryPlannerExtended", "executionStats", "allPlansExecution"])
+            .optional()
+            .default("queryPlanner")
+            .describe(
+                "The verbosity of the explain plan, defaults to queryPlanner. If the user wants to know how fast is a query in execution time, use executionStats. It supports all verbosities as defined in the MongoDB Driver."
+            ),
     };
 
     public operationType: OperationType = "metadata";
 
-    static readonly defaultVerbosity = ExplainVerbosity.queryPlanner;
-
     protected async execute({
         database,
         collection,
         method: methods,
+        verbosity,
     }: ToolArgs<typeof this.argsShape>): Promise<CallToolResult> {
         const provider = await this.ensureConnected();
         const method = methods[0];
@@ -66,14 +71,12 @@ export class ExplainTool extends MongoDBToolBase {
                             writeConcern: undefined,
                         }
                     )
-                    .explain(ExplainTool.defaultVerbosity);
+                    .explain(verbosity);
                 break;
             }
             case "find": {
                 const { filter, ...rest } = method.arguments;
-                result = await provider
-                    .find(database, collection, filter as Document, { ...rest })
-                    .explain(ExplainTool.defaultVerbosity);
+                result = await provider.find(database, collection, filter as Document, { ...rest }).explain(verbosity);
                 break;
             }
             case "count": {
@@ -83,15 +86,15 @@ export class ExplainTool extends MongoDBToolBase {
                         count: collection,
                         query,
                     },
-                    verbosity: ExplainTool.defaultVerbosity,
+                    verbosity,
                 });
                 break;
             }
         }
 
         return {
             content: formatUntrustedData(
-                `Here is some information about the winning plan chosen by the query optimizer for running the given \`${method.name}\` operation in "${database}.${collection}". This information can be used to understand how the query was executed and to optimize the query performance.`,
+                `Here is some information about the winning plan chosen by the query optimizer for running the given \`${method.name}\` operation in "${database}.${collection}". The execution plan was run with the following verbosity: "${verbosity}". This information can be used to understand how the query was executed and to optimize the query performance.`,
                 JSON.stringify(result)
             ),
         };

tests/integration/tools/mongodb/metadata/explain.test.ts

@@ -21,6 +21,13 @@ describeWithMongoDB("explain tool", (integration) => {
                 type: "array",
                 required: true,
             },
+            {
+                name: "verbosity",
+                description:
+                    "The verbosity of the explain plan, defaults to queryPlanner. If the user wants to know how fast is a query in execution time, use executionStats. It supports all verbosities as defined in the MongoDB Driver.",
+                type: "string",
+                required: false,
+            },
         ]
     );
 
@@ -53,7 +60,53 @@ describeWithMongoDB("explain tool", (integration) => {
     for (const testType of ["database", "collection"] as const) {
         describe(`with non-existing ${testType}`, () => {
             for (const testCase of testCases) {
-                it(`should return the explain plan for ${testCase.method}`, async () => {
+                it(`should return the explain plan for "queryPlanner" verbosity for ${testCase.method}`, async () => {
+                    if (testType === "database") {
+                        const { databases } = await integration.mongoClient().db("").admin().listDatabases();
+                        expect(databases.find((db) => db.name === integration.randomDbName())).toBeUndefined();
+                    } else if (testType === "collection") {
+                        await integration
+                            .mongoClient()
+                            .db(integration.randomDbName())
+                            .createCollection("some-collection");
+
+                        const collections = await integration
+                            .mongoClient()
+                            .db(integration.randomDbName())
+                            .listCollections()
+                            .toArray();
+
+                        expect(collections.find((collection) => collection.name === "coll1")).toBeUndefined();
+                    }
+
+                    await integration.connectMcpClient();
+
+                    const response = await integration.mcpClient().callTool({
+                        name: "explain",
+                        arguments: {
+                            database: integration.randomDbName(),
+                            collection: "coll1",
+                            method: [
+                                {
+                                    name: testCase.method,
+                                    arguments: testCase.arguments,
+                                },
+                            ],
+                        },
+                    });
+
+                    const content = getResponseElements(response.content);
+                    expect(content).toHaveLength(2);
+                    expect(content[0]?.text).toEqual(
+                        `Here is some information about the winning plan chosen by the query optimizer for running the given \`${testCase.method}\` operation in "${integration.randomDbName()}.coll1". The execution plan was run with the following verbosity: "queryPlanner". This information can be used to understand how the query was executed and to optimize the query performance.`
+                    );
+
+                    expect(content[1]?.text).toContain("queryPlanner");
+                    expect(content[1]?.text).toContain("winningPlan");
+                    expect(content[1]?.text).not.toContain("executionStats");
+                });
+
+                it(`should return the explain plan for "executionStats" verbosity for ${testCase.method}`, async () => {
                     if (testType === "database") {
                         const { databases } = await integration.mongoClient().db("").admin().listDatabases();
                         expect(databases.find((db) => db.name === integration.randomDbName())).toBeUndefined();
@@ -85,17 +138,19 @@ describeWithMongoDB("explain tool", (integration) => {
                                     arguments: testCase.arguments,
                                 },
                             ],
+                            verbosity: "executionStats",
                         },
                     });
 
                     const content = getResponseElements(response.content);
                     expect(content).toHaveLength(2);
                     expect(content[0]?.text).toEqual(
-                        `Here is some information about the winning plan chosen by the query optimizer for running the given \`${testCase.method}\` operation in "${integration.randomDbName()}.coll1". This information can be used to understand how the query was executed and to optimize the query performance.`
+                        `Here is some information about the winning plan chosen by the query optimizer for running the given \`${testCase.method}\` operation in "${integration.randomDbName()}.coll1". The execution plan was run with the following verbosity: "executionStats". This information can be used to understand how the query was executed and to optimize the query performance.`
                     );
 
                     expect(content[1]?.text).toContain("queryPlanner");
                     expect(content[1]?.text).toContain("winningPlan");
+                    expect(content[1]?.text).toContain("executionStats");
                 });
             }
         });
@@ -121,7 +176,7 @@ describeWithMongoDB("explain tool", (integration) => {
                 });
 
                 for (const testCase of testCases) {
-                    it(`should return the explain plan for ${testCase.method}`, async () => {
+                    it(`should return the explain plan with verbosity "queryPlanner" for ${testCase.method}`, async () => {
                         await integration.connectMcpClient();
 
                         const response = await integration.mcpClient().callTool({
@@ -141,7 +196,7 @@ describeWithMongoDB("explain tool", (integration) => {
                         const content = getResponseElements(response.content);
                         expect(content).toHaveLength(2);
                         expect(content[0]?.text).toEqual(
-                            `Here is some information about the winning plan chosen by the query optimizer for running the given \`${testCase.method}\` operation in "${integration.randomDbName()}.people". This information can be used to understand how the query was executed and to optimize the query performance.`
+                            `Here is some information about the winning plan chosen by the query optimizer for running the given \`${testCase.method}\` operation in "${integration.randomDbName()}.people". The execution plan was run with the following verbosity: "queryPlanner". This information can be used to understand how the query was executed and to optimize the query performance.`
                         );
 
                         expect(content[1]?.text).toContain("queryPlanner");