-
Notifications
You must be signed in to change notification settings - Fork 54
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
base: main
Are you sure you want to change the base?
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -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:[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 | ||
|
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 { | ||||||
const stage = explainResult?.queryPlanner?.winningPlan?.stage; | ||||||
Check failure on line 10 in src/helpers/indexCheck.ts
|
||||||
const inputStage = explainResult?.queryPlanner?.winningPlan?.inputStage; | ||||||
Check failure on line 11 in src/helpers/indexCheck.ts
|
||||||
|
||||||
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
|
||||||
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<Document> | ||||||
): Promise<void> { | ||||||
try { | ||||||
const explainResult = await explainCallback(); | ||||||
|
||||||
if (!usesIndex(explainResult)) { | ||||||
throw new Error(getIndexCheckErrorMessage(database, collection, operation)); | ||||||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
Suggested change
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); | ||||||
} | ||||||
} | ||||||
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"); | ||
}); | ||
}); | ||
}); |
There was a problem hiding this comment.
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.