diff --git a/.github/pull_request_template.md b/.github/pull_request_template.md new file mode 100644 index 00000000..1ebbe374 --- /dev/null +++ b/.github/pull_request_template.md @@ -0,0 +1,5 @@ +## Proposed changes + +## Checklist + +- [ ] I have signed the [MongoDB CLA](https://www.mongodb.com/legal/contributor-agreement) diff --git a/.github/workflows/code_health.yaml b/.github/workflows/code_health.yaml index 2f8ed17a..5505ec67 100644 --- a/.github/workflows/code_health.yaml +++ b/.github/workflows/code_health.yaml @@ -11,7 +11,7 @@ permissions: {} jobs: run-tests: name: Run MongoDB tests - if: github.event.pull_request.user.login != 'dependabot[bot]' && github.event.pull_request.head.repo.full_name == github.repository + if: github.event_name == 'push' || (github.event.pull_request.user.login != 'dependabot[bot]' && github.event.pull_request.head.repo.full_name == github.repository) strategy: matrix: os: [ubuntu-latest, macos-latest, windows-latest] @@ -38,7 +38,7 @@ jobs: run-atlas-tests: name: Run Atlas tests - if: github.event.pull_request.user.login != 'dependabot[bot]' && github.event.pull_request.head.repo.full_name == github.repository + if: github.event_name == 'push' || (github.event.pull_request.user.login != 'dependabot[bot]' && github.event.pull_request.head.repo.full_name == github.repository) runs-on: ubuntu-latest steps: - uses: GitHubSecurityLab/actions-permissions/monitor@v1 @@ -64,7 +64,7 @@ jobs: coverage: name: Report Coverage - if: always() && github.event.pull_request.user.login != 'dependabot[bot]' && github.event.pull_request.head.repo.full_name == github.repository + if: always() && (github.event_name == 'push' || (github.event.pull_request.user.login != 'dependabot[bot]' && github.event.pull_request.head.repo.full_name == github.repository)) runs-on: ubuntu-latest needs: [run-tests, run-atlas-tests] steps: diff --git a/.github/workflows/docker.yaml b/.github/workflows/docker.yaml index 2fd17cf7..ccd07747 100644 --- a/.github/workflows/docker.yaml +++ b/.github/workflows/docker.yaml @@ -16,7 +16,7 @@ jobs: - name: Check out code uses: actions/checkout@b4ffde65f46336ab88eb53be808477a3936bae11 - name: Set up Docker Buildx - uses: docker/setup-buildx-action@18ce135bb5112fa8ce4ed6c17ab05699d7f3a5e0 + uses: docker/setup-buildx-action@e468171a9de216ec08956ac3ada2f0791b6bd435 - name: Login to Docker Hub uses: docker/login-action@74a5d142397b4f367a81961eba4e8cd7edddf772 with: 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..e211feb9 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/package-lock.json b/package-lock.json index 041de7d2..7c3ddb32 100644 --- a/package-lock.json +++ b/package-lock.json @@ -1,12 +1,12 @@ { "name": "mongodb-mcp-server", - "version": "0.1.2", + "version": "0.1.3", "lockfileVersion": 3, "requires": true, "packages": { "": { "name": "mongodb-mcp-server", - "version": "0.1.2", + "version": "0.1.3", "license": "Apache-2.0", "dependencies": { "@modelcontextprotocol/sdk": "^1.11.2", @@ -40,7 +40,7 @@ "@types/yargs-parser": "^21.0.3", "eslint": "^9.24.0", "eslint-config-prettier": "^10.1.2", - "eslint-plugin-jest": "^28.11.0", + "eslint-plugin-jest": "^29.0.1", "eslint-plugin-prettier": "^5.2.6", "globals": "^16.0.0", "jest": "^29.7.0", @@ -1783,9 +1783,9 @@ } }, "node_modules/@eslint/config-array": { - "version": "0.20.0", - "resolved": "https://registry.npmjs.org/@eslint/config-array/-/config-array-0.20.0.tgz", - "integrity": "sha512-fxlS1kkIjx8+vy2SjuCB94q3htSNrufYTXubwiBFeaQHbH6Ipi43gFJq2zCMt6PHhImH3Xmr0NksKDvchWlpQQ==", + "version": "0.20.1", + "resolved": "https://registry.npmjs.org/@eslint/config-array/-/config-array-0.20.1.tgz", + "integrity": "sha512-OL0RJzC/CBzli0DrrR31qzj6d6i6Mm3HByuhflhl4LOBiWxN+3i6/t/ZQQNii4tjksXi8r2CRW1wMpWA2ULUEw==", "dev": true, "license": "Apache-2.0", "dependencies": { @@ -1798,9 +1798,9 @@ } }, "node_modules/@eslint/config-array/node_modules/brace-expansion": { - "version": "1.1.11", - "resolved": "https://registry.npmjs.org/brace-expansion/-/brace-expansion-1.1.11.tgz", - "integrity": "sha512-iCuPHDFgrHX7H2vEI/5xpz07zSHB00TpugqhmYtVmMO6518mCuRMoOYFldEBl0g187ufozdaHgWKcYFb61qGiA==", + "version": "1.1.12", + "resolved": "https://registry.npmjs.org/brace-expansion/-/brace-expansion-1.1.12.tgz", + "integrity": "sha512-9T9UjW3r0UW5c1Q7GTwllptXwhvYmEzFhzMfZ9H7FQWt+uZePjZPjBP/W1ZEyZ1twGWom5/56TF4lPcqjnDHcg==", "dev": true, "license": "MIT", "dependencies": { @@ -1832,9 +1832,9 @@ } }, "node_modules/@eslint/core": { - "version": "0.13.0", - "resolved": "https://registry.npmjs.org/@eslint/core/-/core-0.13.0.tgz", - "integrity": "sha512-yfkgDw1KR66rkT5A8ci4irzDysN7FRpq3ttJolR88OqQikAWqwA8j5VZyas+vjyBNFIJ7MfybJ9plMILI2UrCw==", + "version": "0.14.0", + "resolved": "https://registry.npmjs.org/@eslint/core/-/core-0.14.0.tgz", + "integrity": "sha512-qIbV0/JZr7iSDjqAc60IqbLdsj9GDt16xQtWD+B78d/HAlvysGdZZ6rpJHGAc2T0FQx1X6thsSPdnoiGKdNtdg==", "dev": true, "license": "Apache-2.0", "dependencies": { @@ -1930,9 +1930,9 @@ } }, "node_modules/@eslint/js": { - "version": "9.28.0", - "resolved": "https://registry.npmjs.org/@eslint/js/-/js-9.28.0.tgz", - "integrity": "sha512-fnqSjGWd/CoIp4EXIxWVK/sHA6DOHN4+8Ix2cX5ycOY7LG0UY8nHCU5pIp2eaE1Mc7Qd8kHspYNzYXT2ojPLzg==", + "version": "9.29.0", + "resolved": "https://registry.npmjs.org/@eslint/js/-/js-9.29.0.tgz", + "integrity": "sha512-3PIF4cBw/y+1u2EazflInpV+lYsSG0aByVIQzAgb1m1MhHFSbqTyNqtBKHgWf/9Ykud+DhILS9EGkmekVhbKoQ==", "dev": true, "license": "MIT", "engines": { @@ -1953,19 +1953,32 @@ } }, "node_modules/@eslint/plugin-kit": { - "version": "0.2.8", - "resolved": "https://registry.npmjs.org/@eslint/plugin-kit/-/plugin-kit-0.2.8.tgz", - "integrity": "sha512-ZAoA40rNMPwSm+AeHpCq8STiNAwzWLJuP8Xv4CHIc9wv/PSuExjMrmjfYNj682vW0OOiZ1HKxzvjQr9XZIisQA==", + "version": "0.3.2", + "resolved": "https://registry.npmjs.org/@eslint/plugin-kit/-/plugin-kit-0.3.2.tgz", + "integrity": "sha512-4SaFZCNfJqvk/kenHpI8xvN42DMaoycy4PzKc5otHxRswww1kAt82OlBuwRVLofCACCTZEcla2Ydxv8scMXaTg==", "dev": true, "license": "Apache-2.0", "dependencies": { - "@eslint/core": "^0.13.0", + "@eslint/core": "^0.15.0", "levn": "^0.4.1" }, "engines": { "node": "^18.18.0 || ^20.9.0 || >=21.1.0" } }, + "node_modules/@eslint/plugin-kit/node_modules/@eslint/core": { + "version": "0.15.0", + "resolved": "https://registry.npmjs.org/@eslint/core/-/core-0.15.0.tgz", + "integrity": "sha512-b7ePw78tEWWkpgZCDYkbqDOP8dmM6qe+AOC6iuJqlq1R/0ahMAeH3qynpnqKFGkMltrp44ohV4ubGyvLX28tzw==", + "dev": true, + "license": "Apache-2.0", + "dependencies": { + "@types/json-schema": "^7.0.15" + }, + "engines": { + "node": "^18.18.0 || ^20.9.0 || >=21.1.0" + } + }, "node_modules/@exodus/schemasafe": { "version": "1.3.0", "resolved": "https://registry.npmjs.org/@exodus/schemasafe/-/schemasafe-1.3.0.tgz", @@ -3401,9 +3414,9 @@ } }, "node_modules/@modelcontextprotocol/sdk": { - "version": "1.12.1", - "resolved": "https://registry.npmjs.org/@modelcontextprotocol/sdk/-/sdk-1.12.1.tgz", - "integrity": "sha512-KG1CZhZfWg+u8pxeM/mByJDScJSrjjxLc8fwQqbsS8xCjBmQfMNEBTotYdNanKekepnfRI85GtgQlctLFpcYPw==", + "version": "1.13.1", + "resolved": "https://registry.npmjs.org/@modelcontextprotocol/sdk/-/sdk-1.13.1.tgz", + "integrity": "sha512-8q6+9aF0yA39/qWT/uaIj6zTpC+Qu07DnN/lb9mjoquCJsAh6l3HyYqc9O3t2j7GilseOQOQimLg7W3By6jqvg==", "license": "MIT", "dependencies": { "ajv": "^6.12.6", @@ -6601,9 +6614,9 @@ } }, "node_modules/acorn": { - "version": "8.14.1", - "resolved": "https://registry.npmjs.org/acorn/-/acorn-8.14.1.tgz", - "integrity": "sha512-OvQ/2pUDKmgfCg++xsTX1wGxfTaszcHVcTctW4UJB4hibJx2HXxxO5UmVgyjMa+ZDsiaf5wWLXYpRWMmBI0QHg==", + "version": "8.15.0", + "resolved": "https://registry.npmjs.org/acorn/-/acorn-8.15.0.tgz", + "integrity": "sha512-NZyJarBfL7nWwIq+FDL6Zp/yHEhePMNnnJ0y3qfieCrmNvYct8uvtiV41UvlSe6apAfk0fY1FbWx+NwfmpvtTg==", "dev": true, "license": "MIT", "bin": { @@ -8664,24 +8677,23 @@ } }, "node_modules/eslint": { - "version": "9.26.0", - "resolved": "https://registry.npmjs.org/eslint/-/eslint-9.26.0.tgz", - "integrity": "sha512-Hx0MOjPh6uK9oq9nVsATZKE/Wlbai7KFjfCuw9UHaguDW3x+HF0O5nIi3ud39TWgrTjTO5nHxmL3R1eANinWHQ==", + "version": "9.29.0", + "resolved": "https://registry.npmjs.org/eslint/-/eslint-9.29.0.tgz", + "integrity": "sha512-GsGizj2Y1rCWDu6XoEekL3RLilp0voSePurjZIkxL3wlm5o5EC9VpgaP7lrCvjnkuLvzFBQWB3vWB3K5KQTveQ==", "dev": true, "license": "MIT", "dependencies": { "@eslint-community/eslint-utils": "^4.2.0", "@eslint-community/regexpp": "^4.12.1", - "@eslint/config-array": "^0.20.0", + "@eslint/config-array": "^0.20.1", "@eslint/config-helpers": "^0.2.1", - "@eslint/core": "^0.13.0", + "@eslint/core": "^0.14.0", "@eslint/eslintrc": "^3.3.1", - "@eslint/js": "9.26.0", - "@eslint/plugin-kit": "^0.2.8", + "@eslint/js": "9.29.0", + "@eslint/plugin-kit": "^0.3.1", "@humanfs/node": "^0.16.6", "@humanwhocodes/module-importer": "^1.0.1", "@humanwhocodes/retry": "^0.4.2", - "@modelcontextprotocol/sdk": "^1.8.0", "@types/estree": "^1.0.6", "@types/json-schema": "^7.0.15", "ajv": "^6.12.4", @@ -8689,9 +8701,9 @@ "cross-spawn": "^7.0.6", "debug": "^4.3.2", "escape-string-regexp": "^4.0.0", - "eslint-scope": "^8.3.0", - "eslint-visitor-keys": "^4.2.0", - "espree": "^10.3.0", + "eslint-scope": "^8.4.0", + "eslint-visitor-keys": "^4.2.1", + "espree": "^10.4.0", "esquery": "^1.5.0", "esutils": "^2.0.2", "fast-deep-equal": "^3.1.3", @@ -8705,8 +8717,7 @@ "lodash.merge": "^4.6.2", "minimatch": "^3.1.2", "natural-compare": "^1.4.0", - "optionator": "^0.9.3", - "zod": "^3.24.2" + "optionator": "^0.9.3" }, "bin": { "eslint": "bin/eslint.js" @@ -8743,20 +8754,20 @@ } }, "node_modules/eslint-plugin-jest": { - "version": "28.12.0", - "resolved": "https://registry.npmjs.org/eslint-plugin-jest/-/eslint-plugin-jest-28.12.0.tgz", - "integrity": "sha512-J6zmDp8WiQ9tyvYXE+3RFy7/+l4hraWLzmsabYXyehkmmDd36qV4VQFc7XzcsD8C1PTNt646MSx25bO1mdd9Yw==", + "version": "29.0.1", + "resolved": "https://registry.npmjs.org/eslint-plugin-jest/-/eslint-plugin-jest-29.0.1.tgz", + "integrity": "sha512-EE44T0OSMCeXhDrrdsbKAhprobKkPtJTbQz5yEktysNpHeDZTAL1SfDTNKmcFfJkY6yrQLtTKZALrD3j/Gpmiw==", "dev": true, "license": "MIT", "dependencies": { - "@typescript-eslint/utils": "^6.0.0 || ^7.0.0 || ^8.0.0" + "@typescript-eslint/utils": "^8.0.0" }, "engines": { - "node": "^16.10.0 || ^18.12.0 || >=20.0.0" + "node": "^20.12.0 || ^22.0.0 || >=24.0.0" }, "peerDependencies": { - "@typescript-eslint/eslint-plugin": "^6.0.0 || ^7.0.0 || ^8.0.0", - "eslint": "^7.0.0 || ^8.0.0 || ^9.0.0", + "@typescript-eslint/eslint-plugin": "^8.0.0", + "eslint": "^8.57.0 || ^9.0.0", "jest": "*" }, "peerDependenciesMeta": { @@ -8800,9 +8811,9 @@ } }, "node_modules/eslint-scope": { - "version": "8.3.0", - "resolved": "https://registry.npmjs.org/eslint-scope/-/eslint-scope-8.3.0.tgz", - "integrity": "sha512-pUNxi75F8MJ/GdeKtVLSbYg4ZI34J6C0C7sbL4YOp2exGwen7ZsuBqKzUhXd0qMQ362yET3z+uPwKeg/0C2XCQ==", + "version": "8.4.0", + "resolved": "https://registry.npmjs.org/eslint-scope/-/eslint-scope-8.4.0.tgz", + "integrity": "sha512-sNXOfKCn74rt8RICKMvJS7XKV/Xk9kA7DyJr8mJik3S7Cwgy3qlkkmyS2uQB3jiJg6VNdZd/pDBJu0nvG2NlTg==", "dev": true, "license": "BSD-2-Clause", "dependencies": { @@ -8817,9 +8828,9 @@ } }, "node_modules/eslint-visitor-keys": { - "version": "4.2.0", - "resolved": "https://registry.npmjs.org/eslint-visitor-keys/-/eslint-visitor-keys-4.2.0.tgz", - "integrity": "sha512-UyLnSehNt62FFhSwjZlHmeokpRK59rcz29j+F1/aDgbkbRTk7wIc9XzdoasMUbRNKDM0qQt/+BJ4BrpFeABemw==", + "version": "4.2.1", + "resolved": "https://registry.npmjs.org/eslint-visitor-keys/-/eslint-visitor-keys-4.2.1.tgz", + "integrity": "sha512-Uhdk5sfqcee/9H/rCOJikYz67o0a2Tw2hGRPOG2Y1R2dg7brRe1uG0yaNQDHu+TO/uQPF/5eCapvYSmHUjt7JQ==", "dev": true, "license": "Apache-2.0", "engines": { @@ -8829,16 +8840,6 @@ "url": "https://opencollective.com/eslint" } }, - "node_modules/eslint/node_modules/@eslint/js": { - "version": "9.26.0", - "resolved": "https://registry.npmjs.org/@eslint/js/-/js-9.26.0.tgz", - "integrity": "sha512-I9XlJawFdSMvWjDt6wksMCrgns5ggLNfFwFvnShsleWruvXM514Qxk8V246efTw+eo9JABvVz+u3q2RiAowKxQ==", - "dev": true, - "license": "MIT", - "engines": { - "node": "^18.18.0 || ^20.9.0 || >=21.1.0" - } - }, "node_modules/eslint/node_modules/ajv": { "version": "6.12.6", "resolved": "https://registry.npmjs.org/ajv/-/ajv-6.12.6.tgz", @@ -8901,15 +8902,15 @@ } }, "node_modules/espree": { - "version": "10.3.0", - "resolved": "https://registry.npmjs.org/espree/-/espree-10.3.0.tgz", - "integrity": "sha512-0QYC8b24HWY8zjRnDTL6RiHfDbAWn63qb4LMj1Z4b076A4une81+z03Kg7l7mn/48PUTqoLptSXez8oknU8Clg==", + "version": "10.4.0", + "resolved": "https://registry.npmjs.org/espree/-/espree-10.4.0.tgz", + "integrity": "sha512-j6PAQ2uUr79PZhBjP5C5fhl8e39FmRnOjsD5lGnWrFU8i2G776tBK7+nP8KuQUTTyAZUwfQqXAgrVH5MbH9CYQ==", "dev": true, "license": "BSD-2-Clause", "dependencies": { - "acorn": "^8.14.0", + "acorn": "^8.15.0", "acorn-jsx": "^5.3.2", - "eslint-visitor-keys": "^4.2.0" + "eslint-visitor-keys": "^4.2.1" }, "engines": { "node": "^18.18.0 || ^20.9.0 || >=21.1.0" @@ -11877,9 +11878,9 @@ "optional": true }, "node_modules/mongodb-redact": { - "version": "1.1.6", - "resolved": "https://registry.npmjs.org/mongodb-redact/-/mongodb-redact-1.1.6.tgz", - "integrity": "sha512-L4L3byUH/V/L6YH954NBM/zJpyDHQYmm9eUCxMxqMUfiYCPtmCK1sv/LhxE7UonOkFNEAT6eq2J8gIWGUpHcJA==", + "version": "1.1.8", + "resolved": "https://registry.npmjs.org/mongodb-redact/-/mongodb-redact-1.1.8.tgz", + "integrity": "sha512-EbZ+q7LsVz7q8n49mGIcXgP2UiBp6R6vHEVbmGnF21ThCnP6AIho7wqpHqyjqqGjg54DoXQJTCwHPSknsCHv6g==", "license": "Apache-2.0" }, "node_modules/mongodb-runner": { diff --git a/package.json b/package.json index 54d28a9c..ca262abc 100644 --- a/package.json +++ b/package.json @@ -1,7 +1,7 @@ { "name": "mongodb-mcp-server", "description": "MongoDB Model Context Protocol Server", - "version": "0.1.2", + "version": "0.1.3", "main": "dist/index.js", "author": "MongoDB ", "homepage": "https://github.com/mongodb-js/mongodb-mcp-server", @@ -43,7 +43,7 @@ "@types/yargs-parser": "^21.0.3", "eslint": "^9.24.0", "eslint-config-prettier": "^10.1.2", - "eslint-plugin-jest": "^28.11.0", + "eslint-plugin-jest": "^29.0.1", "eslint-plugin-prettier": "^5.2.6", "globals": "^16.0.0", "jest": "^29.7.0", 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/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 new file mode 100644 index 00000000..22bba447 --- /dev/null +++ b/src/helpers/indexCheck.ts @@ -0,0 +1,83 @@ +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 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 = [ + "IXSCAN", + "COUNT_SCAN", + "EXPRESS_IXSCAN", + "EXPRESS_CLUSTERED_IXSCAN", + "EXPRESS_UPDATE", + "EXPRESS_DELETE", + "IDHACK", + ]; + + if (stage && indexScanStages.includes(stage)) { + return true; + } + + if (inputStage && inputStage.stage && indexScanStages.includes(inputStage.stage as string)) { + 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 MongoDBError( + ErrorCodes.ForbiddenCollscan, + getIndexCheckErrorMessage(database, collection, operation) + ); + } + } catch (error) { + if (error instanceof MongoDBError && error.code === ErrorCodes.ForbiddenCollscan) { + 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); + } +} diff --git a/src/logger.ts b/src/logger.ts index 7adf1263..8157324b 100644 --- a/src/logger.ts +++ b/src/logger.ts @@ -25,7 +25,6 @@ export const LogId = { telemetryMetadataError: mongoLogId(1_002_005), telemetryDeviceIdFailure: mongoLogId(1_002_006), telemetryDeviceIdTimeout: mongoLogId(1_002_007), - telemetryContainerEnvFailure: mongoLogId(1_002_008), toolExecute: mongoLogId(1_003_001), toolExecuteFailure: mongoLogId(1_003_002), diff --git a/src/server.ts b/src/server.ts index 9012fdf5..b0e8e19c 100644 --- a/src/server.ts +++ b/src/server.ts @@ -130,7 +130,7 @@ export class Server { } } - this.telemetry.emitEvents([event]); + this.telemetry.emitEvents([event]).catch(() => {}); } private registerTools() { diff --git a/src/telemetry/telemetry.ts b/src/telemetry/telemetry.ts index 3f543341..ccf0eb41 100644 --- a/src/telemetry/telemetry.ts +++ b/src/telemetry/telemetry.ts @@ -7,152 +7,114 @@ import { MACHINE_METADATA } from "./constants.js"; import { EventCache } from "./eventCache.js"; import nodeMachineId from "node-machine-id"; import { getDeviceId } from "@mongodb-js/device-id"; -import fs from "fs/promises"; - -async function fileExists(filePath: string): Promise { - try { - await fs.access(filePath, fs.constants.F_OK); - return true; // File exists - } catch (e: unknown) { - if ( - e instanceof Error && - ( - e as Error & { - code: string; - } - ).code === "ENOENT" - ) { - return false; // File does not exist - } - throw e; // Re-throw unexpected errors - } -} -async function isContainerized(): Promise { - if (process.env.container) { - return true; - } - - const exists = await Promise.all(["/.dockerenv", "/run/.containerenv", "/var/run/.containerenv"].map(fileExists)); +type EventResult = { + success: boolean; + error?: Error; +}; - return exists.includes(true); -} +export const DEVICE_ID_TIMEOUT = 3000; export class Telemetry { + private isBufferingEvents: boolean = true; + /** Resolves when the device ID is retrieved or timeout occurs */ + public deviceIdPromise: Promise | undefined; private deviceIdAbortController = new AbortController(); private eventCache: EventCache; private getRawMachineId: () => Promise; - private getContainerEnv: () => Promise; - private cachedCommonProperties?: CommonProperties; - private flushing: boolean = false; private constructor( private readonly session: Session, private readonly userConfig: UserConfig, - { - eventCache, - getRawMachineId, - getContainerEnv, - }: { - eventCache: EventCache; - getRawMachineId: () => Promise; - getContainerEnv: () => Promise; - } + private readonly commonProperties: CommonProperties, + { eventCache, getRawMachineId }: { eventCache: EventCache; getRawMachineId: () => Promise } ) { this.eventCache = eventCache; this.getRawMachineId = getRawMachineId; - this.getContainerEnv = getContainerEnv; } static create( session: Session, userConfig: UserConfig, { + commonProperties = { ...MACHINE_METADATA }, eventCache = EventCache.getInstance(), getRawMachineId = () => nodeMachineId.machineId(true), - getContainerEnv = isContainerized, }: { eventCache?: EventCache; getRawMachineId?: () => Promise; - getContainerEnv?: () => Promise; + commonProperties?: CommonProperties; } = {} ): Telemetry { - const instance = new Telemetry(session, userConfig, { - eventCache, - getRawMachineId, - getContainerEnv, - }); + const instance = new Telemetry(session, userConfig, commonProperties, { eventCache, getRawMachineId }); + void instance.start(); return instance; } + private async start(): Promise { + if (!this.isTelemetryEnabled()) { + return; + } + this.deviceIdPromise = getDeviceId({ + getMachineId: () => this.getRawMachineId(), + onError: (reason, error) => { + switch (reason) { + case "resolutionError": + logger.debug(LogId.telemetryDeviceIdFailure, "telemetry", String(error)); + break; + case "timeout": + logger.debug(LogId.telemetryDeviceIdTimeout, "telemetry", "Device ID retrieval timed out"); + break; + case "abort": + // No need to log in the case of aborts + break; + } + }, + abortSignal: this.deviceIdAbortController.signal, + }); + + this.commonProperties.device_id = await this.deviceIdPromise; + + this.isBufferingEvents = false; + } + public async close(): Promise { this.deviceIdAbortController.abort(); - await this.flush(); + this.isBufferingEvents = false; + await this.emitEvents(this.eventCache.getEvents()); } /** * Emits events through the telemetry pipeline * @param events - The events to emit */ - public emitEvents(events: BaseEvent[]): void { - void this.flush(events); + public async emitEvents(events: BaseEvent[]): Promise { + try { + if (!this.isTelemetryEnabled()) { + logger.info(LogId.telemetryEmitFailure, "telemetry", `Telemetry is disabled.`); + return; + } + + await this.emit(events); + } catch { + logger.debug(LogId.telemetryEmitFailure, "telemetry", `Error emitting telemetry events.`); + } } /** * Gets the common properties for events * @returns Object containing common properties for all events */ - private async getCommonProperties(): Promise { - if (!this.cachedCommonProperties) { - let deviceId: string | undefined; - let containerEnv: boolean | undefined; - try { - await Promise.all([ - getDeviceId({ - getMachineId: () => this.getRawMachineId(), - onError: (reason, error) => { - switch (reason) { - case "resolutionError": - logger.debug(LogId.telemetryDeviceIdFailure, "telemetry", String(error)); - break; - case "timeout": - logger.debug( - LogId.telemetryDeviceIdTimeout, - "telemetry", - "Device ID retrieval timed out" - ); - break; - case "abort": - // No need to log in the case of aborts - break; - } - }, - abortSignal: this.deviceIdAbortController.signal, - }).then((id) => { - deviceId = id; - }), - this.getContainerEnv().then((env) => { - containerEnv = env; - }), - ]); - } catch (error: unknown) { - const err = error instanceof Error ? error : new Error(String(error)); - logger.debug(LogId.telemetryDeviceIdFailure, "telemetry", err.message); - } - this.cachedCommonProperties = { - ...MACHINE_METADATA, - mcp_client_version: this.session.agentRunner?.version, - mcp_client_name: this.session.agentRunner?.name, - session_id: this.session.sessionId, - config_atlas_auth: this.session.apiClient.hasCredentials() ? "true" : "false", - config_connection_string: this.userConfig.connectionString ? "true" : "false", - is_container_env: containerEnv ? "true" : "false", - device_id: deviceId, - }; - } - - return this.cachedCommonProperties; + public getCommonProperties(): CommonProperties { + return { + ...this.commonProperties, + mcp_client_version: this.session.agentRunner?.version, + mcp_client_name: this.session.agentRunner?.name, + session_id: this.session.sessionId, + config_atlas_auth: this.session.apiClient.hasCredentials() ? "true" : "false", + config_connection_string: this.userConfig.connectionString ? "true" : "false", + }; } /** @@ -173,74 +135,60 @@ export class Telemetry { } /** - * Attempts to flush events through authenticated and unauthenticated clients + * Attempts to emit events through authenticated and unauthenticated clients * Falls back to caching if both attempts fail */ - public async flush(events?: BaseEvent[]): Promise { - if (!this.isTelemetryEnabled()) { - logger.info(LogId.telemetryEmitFailure, "telemetry", `Telemetry is disabled.`); - return; - } - - if (this.flushing) { - this.eventCache.appendEvents(events ?? []); - process.nextTick(async () => { - // try again if in the middle of a flush - await this.flush(); - }); + private async emit(events: BaseEvent[]): Promise { + if (this.isBufferingEvents) { + this.eventCache.appendEvents(events); return; } - this.flushing = true; + const cachedEvents = this.eventCache.getEvents(); + const allEvents = [...cachedEvents, ...events]; - try { - const cachedEvents = this.eventCache.getEvents(); - const allEvents = [...cachedEvents, ...(events ?? [])]; - if (allEvents.length <= 0) { - this.flushing = false; - return; - } - - logger.debug( - LogId.telemetryEmitStart, - "telemetry", - `Attempting to send ${allEvents.length} events (${cachedEvents.length} cached)` - ); + logger.debug( + LogId.telemetryEmitStart, + "telemetry", + `Attempting to send ${allEvents.length} events (${cachedEvents.length} cached)` + ); - await this.sendEvents(this.session.apiClient, allEvents); + const result = await this.sendEvents(this.session.apiClient, allEvents); + if (result.success) { this.eventCache.clearEvents(); logger.debug( LogId.telemetryEmitSuccess, "telemetry", `Sent ${allEvents.length} events successfully: ${JSON.stringify(allEvents, null, 2)}` ); - } catch (error: unknown) { - logger.debug( - LogId.telemetryEmitFailure, - "telemetry", - `Error sending event to client: ${error instanceof Error ? error.message : String(error)}` - ); - this.eventCache.appendEvents(events ?? []); - process.nextTick(async () => { - // try again - await this.flush(); - }); + return; } - this.flushing = false; + logger.debug( + LogId.telemetryEmitFailure, + "telemetry", + `Error sending event to client: ${result.error instanceof Error ? result.error.message : String(result.error)}` + ); + this.eventCache.appendEvents(events); } /** * Attempts to send events through the provided API client */ - private async sendEvents(client: ApiClient, events: BaseEvent[]): Promise { - const commonProperties = await this.getCommonProperties(); - - await client.sendEvents( - events.map((event) => ({ - ...event, - properties: { ...commonProperties, ...event.properties }, - })) - ); + private async sendEvents(client: ApiClient, events: BaseEvent[]): Promise { + try { + await client.sendEvents( + events.map((event) => ({ + ...event, + properties: { ...this.getCommonProperties(), ...event.properties }, + })) + ); + return { success: true }; + } catch (error) { + return { + success: false, + error: error instanceof Error ? error : new Error(String(error)), + }; + } } } diff --git a/src/telemetry/types.ts b/src/telemetry/types.ts index 05ce8f3f..d77cc010 100644 --- a/src/telemetry/types.ts +++ b/src/telemetry/types.ts @@ -71,5 +71,4 @@ export type CommonProperties = { config_atlas_auth?: TelemetryBoolSet; config_connection_string?: TelemetryBoolSet; session_id?: string; - is_container_env?: TelemetryBoolSet; } & CommonStaticProperties; diff --git a/src/tools/mongodb/delete/deleteMany.ts b/src/tools/mongodb/delete/deleteMany.ts index 6b8351ef..0257d167 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.runCommandWithCheck(database, { + 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/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/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/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..0ed3a192 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.runCommandWithCheck(database, { + 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..97c90e08 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,14 @@ 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..7392135b 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.runCommandWithCheck(database, { + 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/src/tools/tool.ts b/src/tools/tool.ts index 37375f70..b7cce354 100644 --- a/src/tools/tool.ts +++ b/src/tools/tool.ts @@ -74,12 +74,12 @@ export abstract class ToolBase { logger.debug(LogId.toolExecute, "tool", `Executing tool ${this.name}`); const result = await this.execute(...args); - this.emitToolEvent(startTime, result, ...args); + await this.emitToolEvent(startTime, result, ...args).catch(() => {}); return result; } catch (error: unknown) { logger.error(LogId.toolExecuteFailure, "tool", `Error executing ${this.name}: ${error as string}`); const toolResult = await this.handleError(error, args[0] as ToolArgs); - this.emitToolEvent(startTime, toolResult, ...args); + await this.emitToolEvent(startTime, toolResult, ...args).catch(() => {}); return toolResult; } }; @@ -179,11 +179,11 @@ export abstract class ToolBase { * @param result - Whether the command succeeded or failed * @param args - The arguments passed to the tool */ - private emitToolEvent( + private async emitToolEvent( startTime: number, result: CallToolResult, ...args: Parameters> - ): void { + ): Promise { if (!this.telemetry.isTelemetryEnabled()) { return; } @@ -209,6 +209,6 @@ export abstract class ToolBase { event.properties.project_id = metadata.projectId; } - this.telemetry.emitEvents([event]); + await this.telemetry.emitEvents([event]); } } diff --git a/tests/integration/indexCheck.test.ts b/tests/integration/indexCheck.test.ts new file mode 100644 index 00000000..70f53ed5 --- /dev/null +++ b/tests/integration/indexCheck.test.ts @@ -0,0 +1,463 @@ +import { defaultTestConfig, 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 + ], + }, + }); + + 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 + }) + ); + }); +}); diff --git a/tests/integration/server.test.ts b/tests/integration/server.test.ts index 3b4c1858..4c9825f7 100644 --- a/tests/integration/server.test.ts +++ b/tests/integration/server.test.ts @@ -47,11 +47,12 @@ describe("Server integration test", () => { it("should return capabilities", () => { const capabilities = integration.mcpClient().getServerCapabilities(); expectDefined(capabilities); - expect(capabilities.completions).toBeUndefined(); - expect(capabilities.experimental).toBeUndefined(); - expectDefined(capabilities?.tools); expectDefined(capabilities?.logging); - expect(capabilities?.prompts).toBeUndefined(); + expectDefined(capabilities?.completions); + expectDefined(capabilities?.tools); + expectDefined(capabilities?.resources); + expect(capabilities.experimental).toBeUndefined(); + expect(capabilities.prompts).toBeUndefined(); }); }); }); diff --git a/tests/integration/telemetry.test.ts b/tests/integration/telemetry.test.ts new file mode 100644 index 00000000..522c1154 --- /dev/null +++ b/tests/integration/telemetry.test.ts @@ -0,0 +1,28 @@ +import { createHmac } from "crypto"; +import { Telemetry } from "../../src/telemetry/telemetry.js"; +import { Session } from "../../src/session.js"; +import { config } from "../../src/config.js"; +import nodeMachineId from "node-machine-id"; + +describe("Telemetry", () => { + it("should resolve the actual machine ID", async () => { + const actualId: string = await nodeMachineId.machineId(true); + + const actualHashedId = createHmac("sha256", actualId.toUpperCase()).update("atlascli").digest("hex"); + + const telemetry = Telemetry.create( + new Session({ + apiBaseUrl: "", + }), + config + ); + + expect(telemetry.getCommonProperties().device_id).toBe(undefined); + expect(telemetry["isBufferingEvents"]).toBe(true); + + await telemetry.deviceIdPromise; + + expect(telemetry.getCommonProperties().device_id).toBe(actualHashedId); + expect(telemetry["isBufferingEvents"]).toBe(false); + }); +}); diff --git a/tests/integration/tools/mongodb/mongodbHelpers.ts b/tests/integration/tools/mongodb/mongodbHelpers.ts index ca4b09c1..935b27db 100644 --- a/tests/integration/tools/mongodb/mongodbHelpers.ts +++ b/tests/integration/tools/mongodb/mongodbHelpers.ts @@ -70,6 +70,7 @@ export function setupMongoDBIntegrationTest(): MongoDBIntegrationTest { tmpDir: dbsDir, logDir: path.join(tmpDir, "mongodb-runner", "logs"), topology: "standalone", + version: "8.0.10", }); return; diff --git a/tests/unit/indexCheck.test.ts b/tests/unit/indexCheck.test.ts new file mode 100644 index 00000000..82b67e68 --- /dev/null +++ b/tests/unit/indexCheck.test.ts @@ -0,0 +1,149 @@ +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 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: { + 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 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: { + 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"); + }); + }); +}); diff --git a/tests/unit/telemetry.test.ts b/tests/unit/telemetry.test.ts index 3e27f9eb..c1ae28ea 100644 --- a/tests/unit/telemetry.test.ts +++ b/tests/unit/telemetry.test.ts @@ -1,6 +1,6 @@ import { ApiClient } from "../../src/common/atlas/apiClient.js"; import { Session } from "../../src/session.js"; -import { Telemetry } from "../../src/telemetry/telemetry.js"; +import { DEVICE_ID_TIMEOUT, Telemetry } from "../../src/telemetry/telemetry.js"; import { BaseEvent, TelemetryResult } from "../../src/telemetry/types.js"; import { EventCache } from "../../src/telemetry/eventCache.js"; import { config } from "../../src/config.js"; @@ -16,8 +16,6 @@ const MockApiClient = ApiClient as jest.MockedClass; jest.mock("../../src/telemetry/eventCache.js"); const MockEventCache = EventCache as jest.MockedClass; -const nextTick = () => new Promise((resolve) => process.nextTick(resolve)); - describe("Telemetry", () => { const machineId = "test-machine-id"; const hashedMachineId = createHmac("sha256", machineId.toUpperCase()).update("atlascli").digest("hex"); @@ -26,11 +24,6 @@ describe("Telemetry", () => { let mockEventCache: jest.Mocked; let session: Session; let telemetry: Telemetry; - let telemetryConfig: { - eventCache: EventCache; - getRawMachineId: () => Promise; - getContainerEnv: () => Promise; - }; // Helper function to create properly typed test events function createTestEvent(options?: { @@ -84,11 +77,19 @@ describe("Telemetry", () => { expect(appendEvents.length).toBe(appendEventsCalls); if (sendEventsCalledWith) { - expect(sendEvents[0]?.[0]).toMatchObject(sendEventsCalledWith); + expect(sendEvents[0]?.[0]).toEqual( + sendEventsCalledWith.map((event) => ({ + ...event, + properties: { + ...telemetry.getCommonProperties(), + ...event.properties, + }, + })) + ); } if (appendEventsCalledWith) { - expect(appendEvents[0]?.[0]).toMatchObject(appendEventsCalledWith); + expect(appendEvents[0]?.[0]).toEqual(appendEventsCalledWith); } } @@ -124,13 +125,10 @@ describe("Telemetry", () => { setAgentRunner: jest.fn().mockResolvedValue(undefined), } as unknown as Session; - telemetryConfig = { + telemetry = Telemetry.create(session, config, { eventCache: mockEventCache, getRawMachineId: () => Promise.resolve(machineId), - getContainerEnv: () => Promise.resolve(false), - }; - - telemetry = Telemetry.create(session, config, telemetryConfig); + }); config.telemetry = "enabled"; }); @@ -140,8 +138,7 @@ describe("Telemetry", () => { it("should send events successfully", async () => { const testEvent = createTestEvent(); - telemetry.emitEvents([testEvent]); - await nextTick(); // wait for the event to be sent + await telemetry.emitEvents([testEvent]); verifyMockCalls({ sendEventsCalls: 1, @@ -155,8 +152,7 @@ describe("Telemetry", () => { const testEvent = createTestEvent(); - telemetry.emitEvents([testEvent]); - await nextTick(); // wait for the event to be sent + await telemetry.emitEvents([testEvent]); verifyMockCalls({ sendEventsCalls: 1, @@ -179,8 +175,7 @@ describe("Telemetry", () => { // Set up mock to return cached events mockEventCache.getEvents.mockReturnValueOnce([cachedEvent]); - telemetry.emitEvents([newEvent]); - await nextTick(); // wait for the event to be sent + await telemetry.emitEvents([newEvent]); verifyMockCalls({ sendEventsCalls: 1, @@ -189,7 +184,9 @@ describe("Telemetry", () => { }); }); - it("should correctly add common properties to events", async () => { + it("should correctly add common properties to events", () => { + const commonProps = telemetry.getCommonProperties(); + // Use explicit type assertion const expectedProps: Record = { mcp_client_version: "1.0.0", @@ -200,86 +197,48 @@ describe("Telemetry", () => { device_id: hashedMachineId, }; - const testEvent = createTestEvent(); - - telemetry.emitEvents([testEvent]); - await nextTick(); // wait for the event to be sent - - const checkEvent = { - ...testEvent, - properties: { - ...testEvent.properties, - ...expectedProps, - }, - }; - - verifyMockCalls({ - sendEventsCalls: 1, - clearEventsCalls: 1, - sendEventsCalledWith: [checkEvent], - }); + expect(commonProps).toMatchObject(expectedProps); }); - it("should send cache new event while sending another event", async () => { - const newEvent = createTestEvent({ - command: "new-command", - component: "new-component", + describe("machine ID resolution", () => { + beforeEach(() => { + jest.clearAllMocks(); + jest.useFakeTimers(); }); - const newEvent2 = createTestEvent({ - command: "new-command-2", - component: "new-component-2", + afterEach(() => { + jest.clearAllMocks(); + jest.useRealTimers(); }); - telemetry.emitEvents([newEvent]); - telemetry.emitEvents([newEvent2]); + it("should successfully resolve the machine ID", async () => { + telemetry = Telemetry.create(session, config, { + getRawMachineId: () => Promise.resolve(machineId), + }); - await nextTick(); // wait for the event to be sent + expect(telemetry["isBufferingEvents"]).toBe(true); + expect(telemetry.getCommonProperties().device_id).toBe(undefined); - verifyMockCalls({ - sendEventsCalls: 1, - clearEventsCalls: 1, - appendEventsCalls: 1, - sendEventsCalledWith: [newEvent], - appendEventsCalledWith: [newEvent2], - }); - }); + await telemetry.deviceIdPromise; - describe("machine ID resolution", () => { - it("should successfully resolve the machine ID", async () => { - const testEvent = createTestEvent(); - - telemetry.emitEvents([testEvent]); - await nextTick(); // wait for the event to be sent - - const checkEvent = { - ...testEvent, - properties: { - ...testEvent.properties, - device_id: hashedMachineId, - }, - }; - - verifyMockCalls({ - sendEventsCalls: 1, - clearEventsCalls: 1, - sendEventsCalledWith: [checkEvent], - }); + expect(telemetry["isBufferingEvents"]).toBe(false); + expect(telemetry.getCommonProperties().device_id).toBe(hashedMachineId); }); it("should handle machine ID resolution failure", async () => { const loggerSpy = jest.spyOn(logger, "debug"); telemetry = Telemetry.create(session, config, { - ...telemetryConfig, getRawMachineId: () => Promise.reject(new Error("Failed to get device ID")), }); - const testEvent = createTestEvent(); + expect(telemetry["isBufferingEvents"]).toBe(true); + expect(telemetry.getCommonProperties().device_id).toBe(undefined); - telemetry.emitEvents([testEvent]); + await telemetry.deviceIdPromise; - await nextTick(); // wait for the event to be sent + expect(telemetry["isBufferingEvents"]).toBe(false); + expect(telemetry.getCommonProperties().device_id).toBe("unknown"); expect(loggerSpy).toHaveBeenCalledWith( LogId.telemetryDeviceIdFailure, @@ -288,28 +247,27 @@ describe("Telemetry", () => { ); }); - it("should timeout if machine ID resolution takes too long", () => { + it("should timeout if machine ID resolution takes too long", async () => { const loggerSpy = jest.spyOn(logger, "debug"); - jest.useFakeTimers(); - - telemetry = Telemetry.create(session, config, { - ...telemetryConfig, - getRawMachineId: () => new Promise(() => {}), // Never resolves - }); + telemetry = Telemetry.create(session, config, { getRawMachineId: () => new Promise(() => {}) }); - const testEvent = createTestEvent(); + expect(telemetry["isBufferingEvents"]).toBe(true); + expect(telemetry.getCommonProperties().device_id).toBe(undefined); - telemetry.emitEvents([testEvent]); + jest.advanceTimersByTime(DEVICE_ID_TIMEOUT / 2); - jest.advanceTimersByTime(5000); + // Make sure the timeout doesn't happen prematurely. + expect(telemetry["isBufferingEvents"]).toBe(true); + expect(telemetry.getCommonProperties().device_id).toBe(undefined); - jest.useRealTimers(); + jest.advanceTimersByTime(DEVICE_ID_TIMEOUT); - expect(loggerSpy).toHaveBeenCalledTimes(2); + await telemetry.deviceIdPromise; - expect(loggerSpy).toHaveBeenNthCalledWith( - 2, + expect(telemetry.getCommonProperties().device_id).toBe("unknown"); + expect(telemetry["isBufferingEvents"]).toBe(false); + expect(loggerSpy).toHaveBeenCalledWith( LogId.telemetryDeviceIdTimeout, "telemetry", "Device ID retrieval timed out" @@ -330,12 +288,9 @@ describe("Telemetry", () => { it("should not send events", async () => { const testEvent = createTestEvent(); - telemetry.emitEvents([testEvent]); - await nextTick(); // wait for the event to be sent + await telemetry.emitEvents([testEvent]); - verifyMockCalls({ - sendEventsCalls: 0, - }); + verifyMockCalls(); }); }); @@ -358,12 +313,9 @@ describe("Telemetry", () => { it("should not send events", async () => { const testEvent = createTestEvent(); - telemetry.emitEvents([testEvent]); - await nextTick(); // wait for the event to be sent + await telemetry.emitEvents([testEvent]); - verifyMockCalls({ - sendEventsCalls: 0, - }); + verifyMockCalls(); }); }); });