Skip to content

Commit ed60af9

Browse files
kmruizkmruiz
committed
chore: Add support for custom verbosity to the explain plan tool
While queryPlanner is a good default, there are scenarios where we need more verbose outputs to understand why a query is not acting as expected. For example, executionStats can be useful to see rw locks, disk usage and flow throttling. I'm not sure how good are models understanding the whole executionPlan, and maybe we need to remove information that might not be relevant for debugging to reduce the context usage, but this is a decent first approach and we can improve later.
1 parent dd36b1a commit ed60af9

File tree

2 files changed

+71
-13
lines changed

2 files changed

+71
-13
lines changed

‎src/tools/mongodb/metadata/explain.ts‎

Lines changed: 12 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,6 @@ import type { ToolArgs, OperationType } from "../../tool.js";
44
import { formatUntrustedData } from "../../tool.js";
55
import { z } from "zod";
66
import type { Document } from "mongodb";
7-
import { ExplainVerbosity } from "mongodb";
87
import { AggregateArgs } from "../read/aggregate.js";
98
import { FindArgs } from "../read/find.js";
109
import { CountArgs } from "../read/count.js";
@@ -34,16 +33,22 @@ export class ExplainTool extends MongoDBToolBase {
3433
])
3534
)
3635
.describe("The method and its arguments to run"),
36+
verbosity: z
37+
.enum(["queryPlanner", "queryPlannerExtended", "executionStats", "allPlansExecution"])
38+
.optional()
39+
.default("queryPlanner")
40+
.describe(
41+
"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."
42+
),
3743
};
3844

3945
public operationType: OperationType = "metadata";
4046

41-
static readonly defaultVerbosity = ExplainVerbosity.queryPlanner;
42-
4347
protected async execute({
4448
database,
4549
collection,
4650
method: methods,
51+
verbosity,
4752
}: ToolArgs<typeof this.argsShape>): Promise<CallToolResult> {
4853
const provider = await this.ensureConnected();
4954
const method = methods[0];
@@ -66,14 +71,12 @@ export class ExplainTool extends MongoDBToolBase {
6671
writeConcern: undefined,
6772
}
6873
)
69-
.explain(ExplainTool.defaultVerbosity);
74+
.explain(verbosity);
7075
break;
7176
}
7277
case "find": {
7378
const { filter, ...rest } = method.arguments;
74-
result = await provider
75-
.find(database, collection, filter as Document, { ...rest })
76-
.explain(ExplainTool.defaultVerbosity);
79+
result = await provider.find(database, collection, filter as Document, { ...rest }).explain(verbosity);
7780
break;
7881
}
7982
case "count": {
@@ -83,15 +86,15 @@ export class ExplainTool extends MongoDBToolBase {
8386
count: collection,
8487
query,
8588
},
86-
verbosity: ExplainTool.defaultVerbosity,
89+
verbosity,
8790
});
8891
break;
8992
}
9093
}
9194

9295
return {
9396
content: formatUntrustedData(
94-
`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.`,
97+
`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.`,
9598
JSON.stringify(result)
9699
),
97100
};

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

Lines changed: 59 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,13 @@ describeWithMongoDB("explain tool", (integration) => {
2121
type: "array",
2222
required: true,
2323
},
24+
{
25+
name: "verbosity",
26+
description:
27+
"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.",
28+
type: "string",
29+
required: false,
30+
},
2431
]
2532
);
2633

@@ -53,7 +60,53 @@ describeWithMongoDB("explain tool", (integration) => {
5360
for (const testType of ["database", "collection"] as const) {
5461
describe(`with non-existing ${testType}`, () => {
5562
for (const testCase of testCases) {
56-
it(`should return the explain plan for ${testCase.method}`, async () => {
63+
it(`should return the explain plan for "queryPlanner" verbosity for ${testCase.method}`, async () => {
64+
if (testType === "database") {
65+
const { databases } = await integration.mongoClient().db("").admin().listDatabases();
66+
expect(databases.find((db) => db.name === integration.randomDbName())).toBeUndefined();
67+
} else if (testType === "collection") {
68+
await integration
69+
.mongoClient()
70+
.db(integration.randomDbName())
71+
.createCollection("some-collection");
72+
73+
const collections = await integration
74+
.mongoClient()
75+
.db(integration.randomDbName())
76+
.listCollections()
77+
.toArray();
78+
79+
expect(collections.find((collection) => collection.name === "coll1")).toBeUndefined();
80+
}
81+
82+
await integration.connectMcpClient();
83+
84+
const response = await integration.mcpClient().callTool({
85+
name: "explain",
86+
arguments: {
87+
database: integration.randomDbName(),
88+
collection: "coll1",
89+
method: [
90+
{
91+
name: testCase.method,
92+
arguments: testCase.arguments,
93+
},
94+
],
95+
},
96+
});
97+
98+
const content = getResponseElements(response.content);
99+
expect(content).toHaveLength(2);
100+
expect(content[0]?.text).toEqual(
101+
`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.`
102+
);
103+
104+
expect(content[1]?.text).toContain("queryPlanner");
105+
expect(content[1]?.text).toContain("winningPlan");
106+
expect(content[1]?.text).not.toContain("executionStats");
107+
});
108+
109+
it(`should return the explain plan for "executionStats" verbosity for ${testCase.method}`, async () => {
57110
if (testType === "database") {
58111
const { databases } = await integration.mongoClient().db("").admin().listDatabases();
59112
expect(databases.find((db) => db.name === integration.randomDbName())).toBeUndefined();
@@ -85,17 +138,19 @@ describeWithMongoDB("explain tool", (integration) => {
85138
arguments: testCase.arguments,
86139
},
87140
],
141+
verbosity: "executionStats",
88142
},
89143
});
90144

91145
const content = getResponseElements(response.content);
92146
expect(content).toHaveLength(2);
93147
expect(content[0]?.text).toEqual(
94-
`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.`
148+
`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.`
95149
);
96150

97151
expect(content[1]?.text).toContain("queryPlanner");
98152
expect(content[1]?.text).toContain("winningPlan");
153+
expect(content[1]?.text).toContain("executionStats");
99154
});
100155
}
101156
});
@@ -121,7 +176,7 @@ describeWithMongoDB("explain tool", (integration) => {
121176
});
122177

123178
for (const testCase of testCases) {
124-
it(`should return the explain plan for ${testCase.method}`, async () => {
179+
it(`should return the explain plan with verbosity "queryPlanner" for ${testCase.method}`, async () => {
125180
await integration.connectMcpClient();
126181

127182
const response = await integration.mcpClient().callTool({
@@ -141,7 +196,7 @@ describeWithMongoDB("explain tool", (integration) => {
141196
const content = getResponseElements(response.content);
142197
expect(content).toHaveLength(2);
143198
expect(content[0]?.text).toEqual(
144-
`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.`
199+
`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.`
145200
);
146201

147202
expect(content[1]?.text).toContain("queryPlanner");

0 commit comments

Comments
 (0)