From 91e7b9b1885e77dd2122c306d18ea346b7ac2354 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E4=BB=BB=E5=A4=A9=E6=98=93?= Date: Wed, 18 Jun 2025 22:17:10 +0800 Subject: [PATCH 1/7] feat(index-check): add index check functionality before query --- .smithery/smithery.yaml | 10 ++++ README.md | 16 +++++- src/config.ts | 2 + src/helpers/indexCheck.ts | 66 +++++++++++++++++++++ src/tools/mongodb/delete/deleteMany.ts | 20 +++++++ src/tools/mongodb/read/aggregate.ts | 11 ++++ src/tools/mongodb/read/count.ts | 15 +++++ src/tools/mongodb/read/find.ts | 11 ++++ src/tools/mongodb/update/updateMany.ts | 22 +++++++ tests/unit/indexCheck.test.ts | 80 ++++++++++++++++++++++++++ 10 files changed, 252 insertions(+), 1 deletion(-) create mode 100644 src/helpers/indexCheck.ts create mode 100644 tests/unit/indexCheck.test.ts diff --git a/.smithery/smithery.yaml b/.smithery/smithery.yaml index e7de81be..13952c7b 100644 --- a/.smithery/smithery.yaml +++ b/.smithery/smithery.yaml @@ -24,11 +24,17 @@ startCommand: title: Read-only description: When set to true, only allows read and metadata operation types, disabling create/update/delete operations. default: false + indexCheck: + type: boolean + title: Index Check + description: When set to true, enforces that query operations must use an index, rejecting queries that would perform a collection scan. + default: false exampleConfig: atlasClientId: YOUR_ATLAS_CLIENT_ID atlasClientSecret: YOUR_ATLAS_CLIENT_SECRET connectionString: mongodb+srv://USERNAME:PASSWORD@YOUR_CLUSTER.mongodb.net readOnly: true + indexCheck: false commandFunction: # A function that produces the CLI command to start the MCP on stdio. @@ -54,6 +60,10 @@ startCommand: args.push('--connectionString'); args.push(config.connectionString); } + + if (config.indexCheck) { + args.push('--indexCheck'); + } } return { diff --git a/README.md b/README.md index da193cf5..d778a280 100644 --- a/README.md +++ b/README.md @@ -267,6 +267,7 @@ The MongoDB MCP Server can be configured using multiple methods, with the follow | `logPath` | Folder to store logs. | | `disabledTools` | An array of tool names, operation types, and/or categories of tools that will be disabled. | | `readOnly` | When set to true, only allows read and metadata operation types, disabling create/update/delete operations. | +| `indexCheck` | When set to true, enforces that query operations must use an index, rejecting queries that perform a collection scan. | | `telemetry` | When set to disabled, disables telemetry collection. | #### Log Path @@ -312,6 +313,19 @@ You can enable read-only mode using: When read-only mode is active, you'll see a message in the server logs indicating which tools were prevented from registering due to this restriction. +#### Index Check Mode + +The `indexCheck` configuration option allows you to enforce that query operations must use an index. When enabled, queries that perform a collection scan will be rejected to ensure better performance. + +This is useful for scenarios where you want to ensure that database queries are optimized. + +You can enable index check mode using: + +- **Environment variable**: `export MDB_MCP_INDEX_CHECK=true` +- **Command-line argument**: `--indexCheck` + +When index check mode is active, you'll see an error message if a query is rejected due to not using an index. + #### Telemetry The `telemetry` configuration option allows you to disable telemetry collection. When enabled, the MCP server will collect usage data and send it to MongoDB. @@ -430,7 +444,7 @@ export MDB_MCP_LOG_PATH="/path/to/logs" Pass configuration options as command-line arguments when starting the server: ```shell -npx -y mongodb-mcp-server --apiClientId="your-atlas-service-accounts-client-id" --apiClientSecret="your-atlas-service-accounts-client-secret" --connectionString="mongodb+srv://username:password@cluster.mongodb.net/myDatabase" --logPath=/path/to/logs +npx -y mongodb-mcp-server --apiClientId="your-atlas-service-accounts-client-id" --apiClientSecret="your-atlas-service-accounts-client-secret" --connectionString="mongodb+srv://username:password@cluster.mongodb.net/myDatabase" --logPath=/path/to/logs --readOnly --indexCheck ``` #### MCP configuration file examples diff --git a/src/config.ts b/src/config.ts index 9be54452..d9aa0bbc 100644 --- a/src/config.ts +++ b/src/config.ts @@ -23,6 +23,7 @@ export interface UserConfig { connectOptions: ConnectOptions; disabledTools: Array; readOnly?: boolean; + indexCheck?: boolean; } const defaults: UserConfig = { @@ -37,6 +38,7 @@ const defaults: UserConfig = { disabledTools: [], telemetry: "enabled", readOnly: false, + indexCheck: false, }; export const config = { diff --git a/src/helpers/indexCheck.ts b/src/helpers/indexCheck.ts new file mode 100644 index 00000000..7038044f --- /dev/null +++ b/src/helpers/indexCheck.ts @@ -0,0 +1,66 @@ +import { Document } from "mongodb"; +import { NodeDriverServiceProvider } from "@mongosh/service-provider-node-driver"; + +/** + * Check if the query plan uses an index + * @param explainResult The result of the explain query + * @returns true if an index is used, false if it's a full collection scan + */ +export function usesIndex(explainResult: Document): boolean { + const stage = explainResult?.queryPlanner?.winningPlan?.stage; + const inputStage = explainResult?.queryPlanner?.winningPlan?.inputStage; + + if (stage === "IXSCAN" || stage === "COUNT_SCAN") { + return true; + } + + if (inputStage && (inputStage.stage === "IXSCAN" || inputStage.stage === "COUNT_SCAN")) { + return true; + } + + // Recursively check deeper stages + if (inputStage && inputStage.inputStage) { + return usesIndex({ queryPlanner: { winningPlan: inputStage } }); + } + + if (stage === "COLLSCAN") { + return false; + } + + // Default to false (conservative approach) + return false; +} + +/** + * Generate an error message for index check failure + */ +export function getIndexCheckErrorMessage(database: string, collection: string, operation: string): string { + return `Index check failed: The ${operation} operation on "${database}.${collection}" performs a collection scan (COLLSCAN) instead of using an index. Consider adding an index for better performance. Use 'explain' tool for query plan analysis or 'collection-indexes' to view existing indexes. To disable this check, set MDB_MCP_INDEX_CHECK to false.`; +} + +/** + * Generic function to perform index usage check + */ +export async function checkIndexUsage( + provider: NodeDriverServiceProvider, + database: string, + collection: string, + operation: string, + explainCallback: () => Promise +): Promise { + try { + const explainResult = await explainCallback(); + + if (!usesIndex(explainResult)) { + throw new Error(getIndexCheckErrorMessage(database, collection, operation)); + } + } catch (error) { + if (error instanceof Error && error.message.includes("Index check failed")) { + throw error; + } + + // If explain itself fails, log but do not prevent query execution + // This avoids blocking normal queries in special cases (e.g., permission issues) + console.warn(`Index check failed to execute explain for ${operation} on ${database}.${collection}:`, error); + } +} \ No newline at end of file diff --git a/src/tools/mongodb/delete/deleteMany.ts b/src/tools/mongodb/delete/deleteMany.ts index 6b8351ef..e35d3c6b 100644 --- a/src/tools/mongodb/delete/deleteMany.ts +++ b/src/tools/mongodb/delete/deleteMany.ts @@ -2,6 +2,7 @@ import { z } from "zod"; import { CallToolResult } from "@modelcontextprotocol/sdk/types.js"; import { DbOperationArgs, MongoDBToolBase } from "../mongodbTool.js"; import { ToolArgs, OperationType } from "../../tool.js"; +import { checkIndexUsage } from "../../../helpers/indexCheck.js"; export class DeleteManyTool extends MongoDBToolBase { protected name = "delete-many"; @@ -23,6 +24,25 @@ export class DeleteManyTool extends MongoDBToolBase { filter, }: ToolArgs): Promise { const provider = await this.ensureConnected(); + + // Check if delete operation uses an index if enabled + if (this.config.indexCheck) { + await checkIndexUsage(provider, database, collection, "deleteMany", async () => { + return provider.mongoClient.db(database).command({ + explain: { + delete: collection, + deletes: [ + { + q: filter || {}, + limit: 0, // 0 means delete all matching documents + }, + ], + }, + verbosity: "queryPlanner", + }); + }); + } + const result = await provider.deleteMany(database, collection, filter); return { diff --git a/src/tools/mongodb/read/aggregate.ts b/src/tools/mongodb/read/aggregate.ts index c1a46c71..aa21fc5d 100644 --- a/src/tools/mongodb/read/aggregate.ts +++ b/src/tools/mongodb/read/aggregate.ts @@ -3,6 +3,7 @@ import { CallToolResult } from "@modelcontextprotocol/sdk/types.js"; import { DbOperationArgs, MongoDBToolBase } from "../mongodbTool.js"; import { ToolArgs, OperationType } from "../../tool.js"; import { EJSON } from "bson"; +import { checkIndexUsage } from "../../../helpers/indexCheck.js"; export const AggregateArgs = { pipeline: z.array(z.record(z.string(), z.unknown())).describe("An array of aggregation stages to execute"), @@ -23,6 +24,16 @@ export class AggregateTool extends MongoDBToolBase { pipeline, }: ToolArgs): Promise { const provider = await this.ensureConnected(); + + // Check if aggregate operation uses an index if enabled + if (this.config.indexCheck) { + await checkIndexUsage(provider, database, collection, "aggregate", async () => { + return provider + .aggregate(database, collection, pipeline, {}, { writeConcern: undefined }) + .explain("queryPlanner"); + }); + } + const documents = await provider.aggregate(database, collection, pipeline).toArray(); const content: Array<{ text: string; type: "text" }> = [ diff --git a/src/tools/mongodb/read/count.ts b/src/tools/mongodb/read/count.ts index 5d97afa9..c9c79734 100644 --- a/src/tools/mongodb/read/count.ts +++ b/src/tools/mongodb/read/count.ts @@ -2,6 +2,7 @@ import { CallToolResult } from "@modelcontextprotocol/sdk/types.js"; import { DbOperationArgs, MongoDBToolBase } from "../mongodbTool.js"; import { ToolArgs, OperationType } from "../../tool.js"; import { z } from "zod"; +import { checkIndexUsage } from "../../../helpers/indexCheck.js"; export const CountArgs = { query: z @@ -25,6 +26,20 @@ export class CountTool extends MongoDBToolBase { protected async execute({ database, collection, query }: ToolArgs): Promise { const provider = await this.ensureConnected(); + + // Check if count operation uses an index if enabled + if (this.config.indexCheck) { + await checkIndexUsage(provider, database, collection, "count", async () => { + return provider.mongoClient.db(database).command({ + explain: { + count: collection, + query, + }, + verbosity: "queryPlanner", + }); + }); + } + const count = await provider.count(database, collection, query); return { diff --git a/src/tools/mongodb/read/find.ts b/src/tools/mongodb/read/find.ts index c117cf58..214f2e02 100644 --- a/src/tools/mongodb/read/find.ts +++ b/src/tools/mongodb/read/find.ts @@ -4,6 +4,7 @@ import { DbOperationArgs, MongoDBToolBase } from "../mongodbTool.js"; import { ToolArgs, OperationType } from "../../tool.js"; import { SortDirection } from "mongodb"; import { EJSON } from "bson"; +import { checkIndexUsage } from "../../../helpers/indexCheck.js"; export const FindArgs = { filter: z @@ -39,6 +40,16 @@ export class FindTool extends MongoDBToolBase { sort, }: ToolArgs): Promise { const provider = await this.ensureConnected(); + + // Check if find operation uses an index if enabled + if (this.config.indexCheck) { + await checkIndexUsage(provider, database, collection, "find", async () => { + return provider + .find(database, collection, filter, { projection, limit, sort }) + .explain("queryPlanner"); + }); + } + const documents = await provider.find(database, collection, filter, { projection, limit, sort }).toArray(); const content: Array<{ text: string; type: "text" }> = [ diff --git a/src/tools/mongodb/update/updateMany.ts b/src/tools/mongodb/update/updateMany.ts index 187e4633..f64f8549 100644 --- a/src/tools/mongodb/update/updateMany.ts +++ b/src/tools/mongodb/update/updateMany.ts @@ -2,6 +2,7 @@ import { z } from "zod"; import { CallToolResult } from "@modelcontextprotocol/sdk/types.js"; import { DbOperationArgs, MongoDBToolBase } from "../mongodbTool.js"; import { ToolArgs, OperationType } from "../../tool.js"; +import { checkIndexUsage } from "../../../helpers/indexCheck.js"; export class UpdateManyTool extends MongoDBToolBase { protected name = "update-many"; @@ -32,6 +33,27 @@ export class UpdateManyTool extends MongoDBToolBase { upsert, }: ToolArgs): Promise { const provider = await this.ensureConnected(); + + // Check if update operation uses an index if enabled + if (this.config.indexCheck) { + await checkIndexUsage(provider, database, collection, "updateMany", async () => { + return provider.mongoClient.db(database).command({ + explain: { + update: collection, + updates: [ + { + q: filter || {}, + u: update, + upsert: upsert || false, + multi: true, + }, + ], + }, + verbosity: "queryPlanner", + }); + }); + } + const result = await provider.updateMany(database, collection, filter, update, { upsert, }); diff --git a/tests/unit/indexCheck.test.ts b/tests/unit/indexCheck.test.ts new file mode 100644 index 00000000..cc1e7fc3 --- /dev/null +++ b/tests/unit/indexCheck.test.ts @@ -0,0 +1,80 @@ +import { usesIndex, getIndexCheckErrorMessage } from "../../src/helpers/indexCheck.js"; +import { Document } from "mongodb"; + +describe("indexCheck", () => { + describe("usesIndex", () => { + it("should return true for IXSCAN", () => { + const explainResult: Document = { + queryPlanner: { + winningPlan: { + stage: "IXSCAN", + }, + }, + }; + expect(usesIndex(explainResult)).toBe(true); + }); + + it("should return true for COUNT_SCAN", () => { + const explainResult: Document = { + queryPlanner: { + winningPlan: { + stage: "COUNT_SCAN", + }, + }, + }; + expect(usesIndex(explainResult)).toBe(true); + }); + + it("should return false for COLLSCAN", () => { + const explainResult: Document = { + queryPlanner: { + winningPlan: { + stage: "COLLSCAN", + }, + }, + }; + expect(usesIndex(explainResult)).toBe(false); + }); + + it("should return true for nested IXSCAN in inputStage", () => { + const explainResult: Document = { + queryPlanner: { + winningPlan: { + stage: "LIMIT", + inputStage: { + stage: "IXSCAN", + }, + }, + }, + }; + expect(usesIndex(explainResult)).toBe(true); + }); + + it("should return false for unknown stage types", () => { + const explainResult: Document = { + queryPlanner: { + winningPlan: { + stage: "UNKNOWN_STAGE", + }, + }, + }; + expect(usesIndex(explainResult)).toBe(false); + }); + + it("should handle missing queryPlanner", () => { + const explainResult: Document = {}; + expect(usesIndex(explainResult)).toBe(false); + }); + }); + + describe("getIndexCheckErrorMessage", () => { + it("should generate appropriate error message", () => { + const message = getIndexCheckErrorMessage("testdb", "testcoll", "find"); + expect(message).toContain("Index check failed"); + expect(message).toContain("testdb.testcoll"); + expect(message).toContain("find operation"); + expect(message).toContain("collection scan (COLLSCAN)"); + expect(message).toContain("MDB_MCP_INDEX_CHECK"); + }); + }); +}); \ No newline at end of file From 92e68b0ec7565c12fce9f14bb1772e05bf1a0d85 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E4=BB=BB=E5=A4=A9=E6=98=93?= Date: Fri, 20 Jun 2025 22:39:04 +0800 Subject: [PATCH 2/7] feat(errors): enhance index check and proper error handling --- src/errors.ts | 1 + src/helpers/indexCheck.ts | 19 ++++++++-- tests/unit/indexCheck.test.ts | 69 +++++++++++++++++++++++++++++++++++ 3 files changed, 86 insertions(+), 3 deletions(-) diff --git a/src/errors.ts b/src/errors.ts index ae91c3a0..d81867f0 100644 --- a/src/errors.ts +++ b/src/errors.ts @@ -1,6 +1,7 @@ export enum ErrorCodes { NotConnectedToMongoDB = 1_000_000, MisconfiguredConnectionString = 1_000_001, + ForbiddenCollscan = 1_000_002, } export class MongoDBError extends Error { diff --git a/src/helpers/indexCheck.ts b/src/helpers/indexCheck.ts index 7038044f..8721361d 100644 --- a/src/helpers/indexCheck.ts +++ b/src/helpers/indexCheck.ts @@ -1,5 +1,7 @@ import { Document } from "mongodb"; import { NodeDriverServiceProvider } from "@mongosh/service-provider-node-driver"; +import { ErrorCodes, MongoDBError } from "../errors.js"; + /** * Check if the query plan uses an index @@ -10,11 +12,22 @@ export function usesIndex(explainResult: Document): boolean { const stage = explainResult?.queryPlanner?.winningPlan?.stage; const inputStage = explainResult?.queryPlanner?.winningPlan?.inputStage; - if (stage === "IXSCAN" || stage === "COUNT_SCAN") { + // Check for index scan stages (including MongoDB 8.0+ stages) + const indexScanStages = [ + "IXSCAN", + "COUNT_SCAN", + "EXPRESS_IXSCAN", + "EXPRESS_CLUSTERED_IXSCAN", + "EXPRESS_UPDATE", + "EXPRESS_DELETE", + "IDHACK" + ]; + + if (indexScanStages.includes(stage)) { return true; } - if (inputStage && (inputStage.stage === "IXSCAN" || inputStage.stage === "COUNT_SCAN")) { + if (inputStage && indexScanStages.includes(inputStage.stage)) { return true; } @@ -52,7 +65,7 @@ export async function checkIndexUsage( const explainResult = await explainCallback(); if (!usesIndex(explainResult)) { - throw new Error(getIndexCheckErrorMessage(database, collection, operation)); + throw new MongoDBError(ErrorCodes.ForbiddenCollscan, getIndexCheckErrorMessage(database, collection, operation)); } } catch (error) { if (error instanceof Error && error.message.includes("Index check failed")) { diff --git a/tests/unit/indexCheck.test.ts b/tests/unit/indexCheck.test.ts index cc1e7fc3..8a6fbca3 100644 --- a/tests/unit/indexCheck.test.ts +++ b/tests/unit/indexCheck.test.ts @@ -25,6 +25,61 @@ describe("indexCheck", () => { expect(usesIndex(explainResult)).toBe(true); }); + it("should return true for IDHACK", () => { + const explainResult: Document = { + queryPlanner: { + winningPlan: { + stage: "IDHACK", + }, + }, + }; + expect(usesIndex(explainResult)).toBe(true); + }); + + it("should return true for EXPRESS_IXSCAN (MongoDB 8.0+)", () => { + const explainResult: Document = { + queryPlanner: { + winningPlan: { + stage: "EXPRESS_IXSCAN", + }, + }, + }; + expect(usesIndex(explainResult)).toBe(true); + }); + + it("should return true for EXPRESS_CLUSTERED_IXSCAN (MongoDB 8.0+)", () => { + const explainResult: Document = { + queryPlanner: { + winningPlan: { + stage: "EXPRESS_CLUSTERED_IXSCAN", + }, + }, + }; + expect(usesIndex(explainResult)).toBe(true); + }); + + it("should return true for EXPRESS_UPDATE (MongoDB 8.0+)", () => { + const explainResult: Document = { + queryPlanner: { + winningPlan: { + stage: "EXPRESS_UPDATE", + }, + }, + }; + expect(usesIndex(explainResult)).toBe(true); + }); + + it("should return true for EXPRESS_DELETE (MongoDB 8.0+)", () => { + const explainResult: Document = { + queryPlanner: { + winningPlan: { + stage: "EXPRESS_DELETE", + }, + }, + }; + expect(usesIndex(explainResult)).toBe(true); + }); + it("should return false for COLLSCAN", () => { const explainResult: Document = { queryPlanner: { @@ -50,6 +105,20 @@ describe("indexCheck", () => { expect(usesIndex(explainResult)).toBe(true); }); + it("should return true for nested EXPRESS_IXSCAN in inputStage", () => { + const explainResult: Document = { + queryPlanner: { + winningPlan: { + stage: "SORT", + inputStage: { + stage: "EXPRESS_IXSCAN", + }, + }, + }, + }; + expect(usesIndex(explainResult)).toBe(true); + }); + it("should return false for unknown stage types", () => { const explainResult: Document = { queryPlanner: { From e8e306592684f05de68a5d2e3c6f79a60fc67766 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E4=BB=BB=E5=A4=A9=E6=98=93?= Date: Sat, 21 Jun 2025 00:26:50 +0800 Subject: [PATCH 3/7] fix(lint-errors): resolve ESLint formatting and type safety issues --- README.md | 2 +- src/helpers/indexCheck.ts | 20 ++++++++++++-------- src/tools/mongodb/mongodbTool.ts | 10 ++++++++++ src/tools/mongodb/read/find.ts | 4 +--- tests/unit/indexCheck.test.ts | 2 +- 5 files changed, 25 insertions(+), 13 deletions(-) diff --git a/README.md b/README.md index d778a280..e211feb9 100644 --- a/README.md +++ b/README.md @@ -267,7 +267,7 @@ The MongoDB MCP Server can be configured using multiple methods, with the follow | `logPath` | Folder to store logs. | | `disabledTools` | An array of tool names, operation types, and/or categories of tools that will be disabled. | | `readOnly` | When set to true, only allows read and metadata operation types, disabling create/update/delete operations. | -| `indexCheck` | When set to true, enforces that query operations must use an index, rejecting queries that perform a collection scan. | +| `indexCheck` | When set to true, enforces that query operations must use an index, rejecting queries that perform a collection scan. | | `telemetry` | When set to disabled, disables telemetry collection. | #### Log Path diff --git a/src/helpers/indexCheck.ts b/src/helpers/indexCheck.ts index 8721361d..ed4f7a72 100644 --- a/src/helpers/indexCheck.ts +++ b/src/helpers/indexCheck.ts @@ -2,15 +2,16 @@ import { Document } from "mongodb"; import { NodeDriverServiceProvider } from "@mongosh/service-provider-node-driver"; import { ErrorCodes, MongoDBError } from "../errors.js"; - /** * Check if the query plan uses an index * @param explainResult The result of the explain query * @returns true if an index is used, false if it's a full collection scan */ export function usesIndex(explainResult: Document): boolean { - const stage = explainResult?.queryPlanner?.winningPlan?.stage; - const inputStage = explainResult?.queryPlanner?.winningPlan?.inputStage; + const queryPlanner = explainResult?.queryPlanner as Document | undefined; + const winningPlan = queryPlanner?.winningPlan as Document | undefined; + const stage = winningPlan?.stage as string | undefined; + const inputStage = winningPlan?.inputStage as Document | undefined; // Check for index scan stages (including MongoDB 8.0+ stages) const indexScanStages = [ @@ -20,14 +21,14 @@ export function usesIndex(explainResult: Document): boolean { "EXPRESS_CLUSTERED_IXSCAN", "EXPRESS_UPDATE", "EXPRESS_DELETE", - "IDHACK" + "IDHACK", ]; - if (indexScanStages.includes(stage)) { + if (stage && indexScanStages.includes(stage)) { return true; } - if (inputStage && indexScanStages.includes(inputStage.stage)) { + if (inputStage && inputStage.stage && indexScanStages.includes(inputStage.stage as string)) { return true; } @@ -65,7 +66,10 @@ export async function checkIndexUsage( const explainResult = await explainCallback(); if (!usesIndex(explainResult)) { - throw new MongoDBError(ErrorCodes.ForbiddenCollscan, getIndexCheckErrorMessage(database, collection, operation)); + throw new MongoDBError( + ErrorCodes.ForbiddenCollscan, + getIndexCheckErrorMessage(database, collection, operation) + ); } } catch (error) { if (error instanceof Error && error.message.includes("Index check failed")) { @@ -76,4 +80,4 @@ export async function checkIndexUsage( // This avoids blocking normal queries in special cases (e.g., permission issues) console.warn(`Index check failed to execute explain for ${operation} on ${database}.${collection}:`, error); } -} \ No newline at end of file +} diff --git a/src/tools/mongodb/mongodbTool.ts b/src/tools/mongodb/mongodbTool.ts index 2ef1aee0..f215f9a2 100644 --- a/src/tools/mongodb/mongodbTool.ts +++ b/src/tools/mongodb/mongodbTool.ts @@ -64,6 +64,16 @@ export abstract class MongoDBToolBase extends ToolBase { ], isError: true, }; + case ErrorCodes.ForbiddenCollscan: + return { + content: [ + { + type: "text", + text: error.message, + }, + ], + isError: true, + }; } } diff --git a/src/tools/mongodb/read/find.ts b/src/tools/mongodb/read/find.ts index 214f2e02..97c90e08 100644 --- a/src/tools/mongodb/read/find.ts +++ b/src/tools/mongodb/read/find.ts @@ -44,9 +44,7 @@ export class FindTool extends MongoDBToolBase { // Check if find operation uses an index if enabled if (this.config.indexCheck) { await checkIndexUsage(provider, database, collection, "find", async () => { - return provider - .find(database, collection, filter, { projection, limit, sort }) - .explain("queryPlanner"); + return provider.find(database, collection, filter, { projection, limit, sort }).explain("queryPlanner"); }); } diff --git a/tests/unit/indexCheck.test.ts b/tests/unit/indexCheck.test.ts index 8a6fbca3..82b67e68 100644 --- a/tests/unit/indexCheck.test.ts +++ b/tests/unit/indexCheck.test.ts @@ -146,4 +146,4 @@ describe("indexCheck", () => { expect(message).toContain("MDB_MCP_INDEX_CHECK"); }); }); -}); \ No newline at end of file +}); From 2dd9f95fe110f327faf887e35f17fd434613a826 Mon Sep 17 00:00:00 2001 From: David Chen <73060999+Crushdada@users.noreply.github.com> Date: Sat, 21 Jun 2025 00:30:50 +0800 Subject: [PATCH 4/7] Update src/helpers/indexCheck.ts Co-authored-by: Himanshu Singh --- src/helpers/indexCheck.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/helpers/indexCheck.ts b/src/helpers/indexCheck.ts index ed4f7a72..22bba447 100644 --- a/src/helpers/indexCheck.ts +++ b/src/helpers/indexCheck.ts @@ -72,7 +72,7 @@ export async function checkIndexUsage( ); } } catch (error) { - if (error instanceof Error && error.message.includes("Index check failed")) { + if (error instanceof MongoDBError && error.code === ErrorCodes.ForbiddenCollscan) { throw error; } From af5e40184d49f496276a29d37200d4f068c8abae Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E4=BB=BB=E5=A4=A9=E6=98=93?= Date: Sat, 21 Jun 2025 00:44:10 +0800 Subject: [PATCH 5/7] fix: use runCommandWithCheck for database commands --- src/tools/mongodb/delete/deleteMany.ts | 2 +- src/tools/mongodb/metadata/explain.ts | 2 +- src/tools/mongodb/read/count.ts | 2 +- src/tools/mongodb/update/updateMany.ts | 2 +- 4 files changed, 4 insertions(+), 4 deletions(-) diff --git a/src/tools/mongodb/delete/deleteMany.ts b/src/tools/mongodb/delete/deleteMany.ts index e35d3c6b..0257d167 100644 --- a/src/tools/mongodb/delete/deleteMany.ts +++ b/src/tools/mongodb/delete/deleteMany.ts @@ -28,7 +28,7 @@ export class DeleteManyTool extends MongoDBToolBase { // Check if delete operation uses an index if enabled if (this.config.indexCheck) { await checkIndexUsage(provider, database, collection, "deleteMany", async () => { - return provider.mongoClient.db(database).command({ + return provider.runCommandWithCheck(database, { explain: { delete: collection, deletes: [ diff --git a/src/tools/mongodb/metadata/explain.ts b/src/tools/mongodb/metadata/explain.ts index e529e899..1068a008 100644 --- a/src/tools/mongodb/metadata/explain.ts +++ b/src/tools/mongodb/metadata/explain.ts @@ -76,7 +76,7 @@ export class ExplainTool extends MongoDBToolBase { } case "count": { const { query } = method.arguments; - result = await provider.mongoClient.db(database).command({ + result = await provider.runCommandWithCheck(database, { explain: { count: collection, query, diff --git a/src/tools/mongodb/read/count.ts b/src/tools/mongodb/read/count.ts index c9c79734..0ed3a192 100644 --- a/src/tools/mongodb/read/count.ts +++ b/src/tools/mongodb/read/count.ts @@ -30,7 +30,7 @@ export class CountTool extends MongoDBToolBase { // Check if count operation uses an index if enabled if (this.config.indexCheck) { await checkIndexUsage(provider, database, collection, "count", async () => { - return provider.mongoClient.db(database).command({ + return provider.runCommandWithCheck(database, { explain: { count: collection, query, diff --git a/src/tools/mongodb/update/updateMany.ts b/src/tools/mongodb/update/updateMany.ts index f64f8549..7392135b 100644 --- a/src/tools/mongodb/update/updateMany.ts +++ b/src/tools/mongodb/update/updateMany.ts @@ -37,7 +37,7 @@ export class UpdateManyTool extends MongoDBToolBase { // Check if update operation uses an index if enabled if (this.config.indexCheck) { await checkIndexUsage(provider, database, collection, "updateMany", async () => { - return provider.mongoClient.db(database).command({ + return provider.runCommandWithCheck(database, { explain: { update: collection, updates: [ From 39337101613d9ba4423d523a3fe210cabcc8983d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E4=BB=BB=E5=A4=A9=E6=98=93?= Date: Thu, 26 Jun 2025 01:43:27 +0800 Subject: [PATCH 6/7] feat: add integration tests for indexCheck config --- tests/integration/indexCheck.test.ts | 464 +++++++++++++++++++++++++++ 1 file changed, 464 insertions(+) create mode 100644 tests/integration/indexCheck.test.ts diff --git a/tests/integration/indexCheck.test.ts b/tests/integration/indexCheck.test.ts new file mode 100644 index 00000000..4aedce2b --- /dev/null +++ b/tests/integration/indexCheck.test.ts @@ -0,0 +1,464 @@ +import { defaultTestConfig, expectDefined, getResponseContent } from "./helpers.js"; +import { describeWithMongoDB } from "./tools/mongodb/mongodbHelpers.js"; + +describe("IndexCheck integration tests", () => { + describe("with indexCheck enabled", () => { + describeWithMongoDB( + "indexCheck functionality", + (integration) => { + beforeEach(async () => { + await integration.connectMcpClient(); + }); + + describe("find operations", () => { + beforeEach(async () => { + // Insert test data for find operations + await integration + .mongoClient() + .db(integration.randomDbName()) + .collection("find-test-collection") + .insertMany([ + { name: "document1", value: 1, category: "A" }, + { name: "document2", value: 2, category: "B" }, + { name: "document3", value: 3, category: "A" }, + ]); + }); + + it("should reject queries that perform collection scans", async () => { + const response = await integration.mcpClient().callTool({ + name: "find", + arguments: { + database: integration.randomDbName(), + collection: "find-test-collection", + filter: { category: "A" }, // No index on category field + }, + }); + + const content = getResponseContent(response.content); + expect(content).toContain("Index check failed"); + expect(content).toContain("collection scan (COLLSCAN)"); + expect(content).toContain("MDB_MCP_INDEX_CHECK"); + expect(response.isError).toBe(true); + }); + + it("should allow queries that use indexes", async () => { + // Create an index on the category field + await integration + .mongoClient() + .db(integration.randomDbName()) + .collection("find-test-collection") + .createIndex({ category: 1 }); + + const response = await integration.mcpClient().callTool({ + name: "find", + arguments: { + database: integration.randomDbName(), + collection: "find-test-collection", + filter: { category: "A" }, // Now has index + }, + }); + + expect(response.isError).toBeFalsy(); + const content = getResponseContent(response.content); + expect(content).toContain("Found"); + expect(content).toContain("documents"); + }); + + it("should allow queries using _id (IDHACK)", async () => { + const docs = await integration + .mongoClient() + .db(integration.randomDbName()) + .collection("find-test-collection") + .find({}) + .toArray(); + + expect(docs.length).toBeGreaterThan(0); + + const response = await integration.mcpClient().callTool({ + name: "find", + arguments: { + database: integration.randomDbName(), + collection: "find-test-collection", + filter: { _id: docs[0]!._id }, // Uses _id index (IDHACK) + }, + }); + + expect(response.isError).toBeFalsy(); + const content = getResponseContent(response.content); + expect(content).toContain("Found 1 documents"); + }); + }); + + describe("count operations", () => { + beforeEach(async () => { + // Insert test data for count operations + await integration + .mongoClient() + .db(integration.randomDbName()) + .collection("count-test-collection") + .insertMany([ + { name: "document1", value: 1, category: "A" }, + { name: "document2", value: 2, category: "B" }, + { name: "document3", value: 3, category: "A" }, + ]); + }); + + it("should reject count queries that perform collection scans", async () => { + const response = await integration.mcpClient().callTool({ + name: "count", + arguments: { + database: integration.randomDbName(), + collection: "count-test-collection", + query: { value: { $gt: 1 } }, // No index on value field + }, + }); + + const content = getResponseContent(response.content); + expect(content).toContain("Index check failed"); + expect(content).toContain("count operation"); + expect(response.isError).toBe(true); + }); + + it("should allow count queries with indexes", async () => { + // Create an index on the value field + await integration + .mongoClient() + .db(integration.randomDbName()) + .collection("count-test-collection") + .createIndex({ value: 1 }); + + const response = await integration.mcpClient().callTool({ + name: "count", + arguments: { + database: integration.randomDbName(), + collection: "count-test-collection", + query: { value: { $gt: 1 } }, // Now has index + }, + }); + + expect(response.isError).toBeFalsy(); + const content = getResponseContent(response.content); + expect(content).toContain("Found"); + expect(content).toMatch(/\d+ documents/); + }); + }); + + describe("aggregate operations", () => { + beforeEach(async () => { + // Insert test data for aggregate operations + await integration + .mongoClient() + .db(integration.randomDbName()) + .collection("aggregate-test-collection") + .insertMany([ + { name: "document1", value: 1, category: "A" }, + { name: "document2", value: 2, category: "B" }, + { name: "document3", value: 3, category: "A" }, + ]); + }); + + it("should reject aggregation queries that perform collection scans", async () => { + const response = await integration.mcpClient().callTool({ + name: "aggregate", + arguments: { + database: integration.randomDbName(), + collection: "aggregate-test-collection", + pipeline: [ + { $match: { category: "A" } }, // No index on category + { $group: { _id: "$category", count: { $sum: 1 } } }, + ], + }, + }); + + const content = getResponseContent(response.content); + expect(content).toContain("Index check failed"); + expect(content).toContain("aggregate operation"); + expect(response.isError).toBe(true); + }); + + it("should allow aggregation queries with indexes", async () => { + // Create an index on the category field + await integration + .mongoClient() + .db(integration.randomDbName()) + .collection("aggregate-test-collection") + .createIndex({ category: 1 }); + + const response = await integration.mcpClient().callTool({ + name: "aggregate", + arguments: { + database: integration.randomDbName(), + collection: "aggregate-test-collection", + pipeline: [ + { $match: { category: "A" } }, // Now has index + { $group: { _id: "$category", count: { $sum: 1 } } }, + ], + }, + }); + + expect(response.isError).toBeFalsy(); + const content = getResponseContent(response.content); + expect(content).toContain("Found"); + }); + }); + + describe("updateMany operations", () => { + beforeEach(async () => { + // Insert test data for updateMany operations + await integration + .mongoClient() + .db(integration.randomDbName()) + .collection("update-test-collection") + .insertMany([ + { name: "document1", value: 1, category: "A" }, + { name: "document2", value: 2, category: "B" }, + { name: "document3", value: 3, category: "A" }, + ]); + }); + + it("should reject updateMany queries that perform collection scans", async () => { + const response = await integration.mcpClient().callTool({ + name: "update-many", + arguments: { + database: integration.randomDbName(), + collection: "update-test-collection", + filter: { category: "A" }, // No index on category + update: { $set: { updated: true } }, + }, + }); + + const content = getResponseContent(response.content); + expect(content).toContain("Index check failed"); + expect(content).toContain("updateMany operation"); + expect(response.isError).toBe(true); + }); + + it("should allow updateMany queries with indexes", async () => { + // Create an index on the category field + await integration + .mongoClient() + .db(integration.randomDbName()) + .collection("update-test-collection") + .createIndex({ category: 1 }); + + const response = await integration.mcpClient().callTool({ + name: "update-many", + arguments: { + database: integration.randomDbName(), + collection: "update-test-collection", + filter: { category: "A" }, // Now has index + update: { $set: { updated: true } }, + }, + }); + + expect(response.isError).toBeFalsy(); + const content = getResponseContent(response.content); + expect(content).toContain("Matched"); + expect(content).toContain("Modified"); + }); + }); + + describe("deleteMany operations", () => { + beforeEach(async () => { + // Insert test data for deleteMany operations + await integration + .mongoClient() + .db(integration.randomDbName()) + .collection("delete-test-collection") + .insertMany([ + { name: "document1", value: 1, category: "A" }, + { name: "document2", value: 2, category: "B" }, + { name: "document3", value: 3, category: "A" }, + ]); + }); + + it("should reject deleteMany queries that perform collection scans", async () => { + const response = await integration.mcpClient().callTool({ + name: "delete-many", + arguments: { + database: integration.randomDbName(), + collection: "delete-test-collection", + filter: { value: { $lt: 2 } }, // No index on value + }, + }); + + const content = getResponseContent(response.content); + expect(content).toContain("Index check failed"); + expect(content).toContain("deleteMany operation"); + expect(response.isError).toBe(true); + }); + + it("should allow deleteMany queries with indexes", async () => { + // Create an index on the value field + await integration + .mongoClient() + .db(integration.randomDbName()) + .collection("delete-test-collection") + .createIndex({ value: 1 }); + + const response = await integration.mcpClient().callTool({ + name: "delete-many", + arguments: { + database: integration.randomDbName(), + collection: "delete-test-collection", + filter: { value: { $lt: 2 } }, // Now has index + }, + }); + + expect(response.isError).toBeFalsy(); + const content = getResponseContent(response.content); + expect(content).toContain("Deleted"); + expect(content).toMatch(/\d+ document\(s\)/); + }); + }); + }, + () => ({ + ...defaultTestConfig, + indexCheck: true, // Enable indexCheck + }) + ); + }); + + describe("with indexCheck disabled", () => { + describeWithMongoDB( + "indexCheck disabled functionality", + (integration) => { + beforeEach(async () => { + await integration.connectMcpClient(); + + // insert test data for disabled indexCheck tests + await integration + .mongoClient() + .db(integration.randomDbName()) + .collection("disabled-test-collection") + .insertMany([ + { name: "document1", value: 1, category: "A" }, + { name: "document2", value: 2, category: "B" }, + { name: "document3", value: 3, category: "A" }, + ]); + }); + + it("should allow all queries regardless of index usage", async () => { + // Test find operation without index + const findResponse = await integration.mcpClient().callTool({ + name: "find", + arguments: { + database: integration.randomDbName(), + collection: "disabled-test-collection", + filter: { category: "A" }, // No index, but should be allowed + }, + }); + + expect(findResponse.isError).toBeFalsy(); + const findContent = getResponseContent(findResponse.content); + expect(findContent).toContain("Found"); + expect(findContent).not.toContain("Index check failed"); + }); + + it("should allow count operations without indexes", async () => { + const response = await integration.mcpClient().callTool({ + name: "count", + arguments: { + database: integration.randomDbName(), + collection: "disabled-test-collection", + query: { value: { $gt: 1 } }, // No index, but should be allowed + }, + }); + + expect(response.isError).toBeFalsy(); + const content = getResponseContent(response.content); + expect(content).toContain("Found"); + expect(content).not.toContain("Index check failed"); + }); + + it("should allow aggregate operations without indexes", async () => { + const response = await integration.mcpClient().callTool({ + name: "aggregate", + arguments: { + database: integration.randomDbName(), + collection: "disabled-test-collection", + pipeline: [ + { $match: { category: "A" } }, // No index, but should be allowed + { $group: { _id: "$category", count: { $sum: 1 } } }, + ], + }, + }); + + expect(response.isError).toBeFalsy(); + const content = getResponseContent(response.content); + expect(content).toContain("Found"); + expect(content).not.toContain("Index check failed"); + }); + + it("should allow updateMany operations without indexes", async () => { + const response = await integration.mcpClient().callTool({ + name: "update-many", + arguments: { + database: integration.randomDbName(), + collection: "disabled-test-collection", + filter: { category: "A" }, // No index, but should be allowed + update: { $set: { updated: true } }, + }, + }); + + expect(response.isError).toBeFalsy(); + const content = getResponseContent(response.content); + expect(content).toContain("Matched"); + expect(content).not.toContain("Index check failed"); + }); + + it("should allow deleteMany operations without indexes", async () => { + const response = await integration.mcpClient().callTool({ + name: "delete-many", + arguments: { + database: integration.randomDbName(), + collection: "disabled-test-collection", + filter: { value: { $lt: 2 } }, // No index, but should be allowed + }, + }); + + expect(response.isError).toBeFalsy(); + const content = getResponseContent(response.content); + expect(content).toContain("Deleted"); + expect(content).not.toContain("Index check failed"); + }); + }, + () => ({ + ...defaultTestConfig, + indexCheck: false, // Disable indexCheck + }) + ); + }); + + describe("indexCheck configuration validation", () => { + describeWithMongoDB( + "default indexCheck behavior", + (integration) => { + it("should allow collection scans by default when indexCheck is not specified", async () => { + await integration.connectMcpClient(); + + await integration + .mongoClient() + .db(integration.randomDbName()) + .collection("default-test-collection") + .insertOne({ name: "test", value: 1 }); + + const response = await integration.mcpClient().callTool({ + name: "find", + arguments: { + database: integration.randomDbName(), + collection: "default-test-collection", + filter: { name: "test" }, // No index, should be allowed by default + }, + }); + + expect(response.isError).toBeFalsy(); + }); + }, + () => ({ + ...defaultTestConfig, + // indexCheck not specified, should default to false + }) + ); + }); +}); \ No newline at end of file From 16a9130c182ec8d18dd2e51bb5f2788f024bc8e8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E4=BB=BB=E5=A4=A9=E6=98=93?= Date: Thu, 26 Jun 2025 17:03:35 +0800 Subject: [PATCH 7/7] fix: code style check failed --- tests/integration/indexCheck.test.ts | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/integration/indexCheck.test.ts b/tests/integration/indexCheck.test.ts index 4aedce2b..c6cc003e 100644 --- a/tests/integration/indexCheck.test.ts +++ b/tests/integration/indexCheck.test.ts @@ -1,4 +1,4 @@ -import { defaultTestConfig, expectDefined, getResponseContent } from "./helpers.js"; +import { defaultTestConfig, getResponseContent } from "./helpers.js"; import { describeWithMongoDB } from "./tools/mongodb/mongodbHelpers.js"; describe("IndexCheck integration tests", () => { @@ -79,7 +79,7 @@ describe("IndexCheck integration tests", () => { arguments: { database: integration.randomDbName(), collection: "find-test-collection", - filter: { _id: docs[0]!._id }, // Uses _id index (IDHACK) + filter: { _id: docs[0]?._id }, // Uses _id index (IDHACK) }, }); @@ -461,4 +461,4 @@ describe("IndexCheck integration tests", () => { }) ); }); -}); \ No newline at end of file +});