Skip to content

feat(index-check): add index check functionality before query #309

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 10 additions & 0 deletions .smithery/smithery.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -54,6 +60,10 @@ startCommand:
args.push('--connectionString');
args.push(config.connectionString);
}

if (config.indexCheck) {
args.push('--indexCheck');
}
}

return {
Expand Down
16 changes: 15 additions & 1 deletion README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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:[email protected]/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:[email protected]/myDatabase" --logPath=/path/to/logs --readOnly --indexCheck
```

#### MCP configuration file examples
Expand Down
2 changes: 2 additions & 0 deletions src/config.ts
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@ export interface UserConfig {
connectOptions: ConnectOptions;
disabledTools: Array<string>;
readOnly?: boolean;
indexCheck?: boolean;
}

const defaults: UserConfig = {
Expand All @@ -37,6 +38,7 @@ const defaults: UserConfig = {
disabledTools: [],
telemetry: "enabled",
readOnly: false,
indexCheck: false,
};

export const config = {
Expand Down
66 changes: 66 additions & 0 deletions src/helpers/indexCheck.ts
Original file line number Diff line number Diff line change
@@ -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 {
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

There will be situations where this doesn't work in MongoDB 8.0+, as we have a few more stages that are like an index scan. Can you add them? These are:

"EXPRESS_IXSCAN"
"EXPRESS_CLUSTERED_IXSCAN"
"EXPRESS_UPDATE"
"EXPRESS_DELETE"

Also, not related to 8.0, but there is a stage called IDHACK that is also like an index scan.

const stage = explainResult?.queryPlanner?.winningPlan?.stage;

Check failure on line 10 in src/helpers/indexCheck.ts

View workflow job for this annotation

GitHub Actions / check-style

Unsafe member access .winningPlan on an `any` value

Check failure on line 10 in src/helpers/indexCheck.ts

View workflow job for this annotation

GitHub Actions / check-style

Unsafe assignment of an `any` value
const inputStage = explainResult?.queryPlanner?.winningPlan?.inputStage;

Check failure on line 11 in src/helpers/indexCheck.ts

View workflow job for this annotation

GitHub Actions / check-style

Unsafe member access .winningPlan on an `any` value

Check failure on line 11 in src/helpers/indexCheck.ts

View workflow job for this annotation

GitHub Actions / check-style

Unsafe assignment of an `any` value

if (stage === "IXSCAN" || stage === "COUNT_SCAN") {
return true;
}

if (inputStage && (inputStage.stage === "IXSCAN" || inputStage.stage === "COUNT_SCAN")) {

Check failure on line 17 in src/helpers/indexCheck.ts

View workflow job for this annotation

GitHub Actions / check-style

Unsafe member access .stage on an `any` value

Check failure on line 17 in src/helpers/indexCheck.ts

View workflow job for this annotation

GitHub Actions / check-style

Unsafe member access .stage on an `any` value
return true;
}

// Recursively check deeper stages
if (inputStage && inputStage.inputStage) {

Check failure on line 22 in src/helpers/indexCheck.ts

View workflow job for this annotation

GitHub Actions / check-style

Unsafe member access .inputStage on an `any` value
return usesIndex({ queryPlanner: { winningPlan: inputStage } });

Check failure on line 23 in src/helpers/indexCheck.ts

View workflow job for this annotation

GitHub Actions / check-style

Unsafe assignment of an `any` value
}

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<Document>
): Promise<void> {
try {
const explainResult = await explainCallback();

if (!usesIndex(explainResult)) {
throw new Error(getIndexCheckErrorMessage(database, collection, operation));
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
throw new Error(getIndexCheckErrorMessage(database, collection, operation));
throw new MongoDBError(ErrorCodes.ForbiddenCollscan, getIndexCheckErrorMessage(database, collection, operation));

Here it's better to use MongoDBError, this way you can rethrow it without taking a look at the message, and it's also what we use as a facade with the MCP Client, so it should be safe to throw anytime!

You'll need to add the error code in src/errors.ts

}
} 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);
}
}

Check failure on line 66 in src/helpers/indexCheck.ts

View workflow job for this annotation

GitHub Actions / check-style

Insert `⏎`
20 changes: 20 additions & 0 deletions src/tools/mongodb/delete/deleteMany.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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";
Expand All @@ -23,6 +24,25 @@ export class DeleteManyTool extends MongoDBToolBase {
filter,
}: ToolArgs<typeof this.argsShape>): Promise<CallToolResult> {
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 {
Expand Down
11 changes: 11 additions & 0 deletions src/tools/mongodb/read/aggregate.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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"),
Expand All @@ -23,6 +24,16 @@ export class AggregateTool extends MongoDBToolBase {
pipeline,
}: ToolArgs<typeof this.argsShape>): Promise<CallToolResult> {
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" }> = [
Expand Down
15 changes: 15 additions & 0 deletions src/tools/mongodb/read/count.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -25,6 +26,20 @@ export class CountTool extends MongoDBToolBase {

protected async execute({ database, collection, query }: ToolArgs<typeof this.argsShape>): Promise<CallToolResult> {
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 {
Expand Down
11 changes: 11 additions & 0 deletions src/tools/mongodb/read/find.ts
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
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
Expand Down Expand Up @@ -39,6 +40,16 @@
sort,
}: ToolArgs<typeof this.argsShape>): Promise<CallToolResult> {
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

Check failure on line 47 in src/tools/mongodb/read/find.ts

View workflow job for this annotation

GitHub Actions / check-style

Replace `⏎····················.find(database,·collection,·filter,·{·projection,·limit,·sort·})⏎····················` with `.find(database,·collection,·filter,·{·projection,·limit,·sort·})`
.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" }> = [
Expand Down
22 changes: 22 additions & 0 deletions src/tools/mongodb/update/updateMany.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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";
Expand Down Expand Up @@ -32,6 +33,27 @@ export class UpdateManyTool extends MongoDBToolBase {
upsert,
}: ToolArgs<typeof this.argsShape>): Promise<CallToolResult> {
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,
});
Expand Down
80 changes: 80 additions & 0 deletions tests/unit/indexCheck.test.ts
Original file line number Diff line number Diff line change
@@ -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");
});
});
});
Loading