From 293cfe2c9dc368d22df64e4a90f501c2495c23a1 Mon Sep 17 00:00:00 2001 From: Himanshu Singh Date: Wed, 20 Aug 2025 10:42:49 +0200 Subject: [PATCH 1/6] feat: adds support for exporting aggregations --- src/common/exportsManager.ts | 6 +- src/tools/mongodb/read/export.ts | 57 +++++-- .../resources/exportedData.test.ts | 21 ++- .../tools/mongodb/read/export.test.ts | 153 ++++++++++++++---- 4 files changed, 182 insertions(+), 55 deletions(-) diff --git a/src/common/exportsManager.ts b/src/common/exportsManager.ts index 384203092..9235e8b25 100644 --- a/src/common/exportsManager.ts +++ b/src/common/exportsManager.ts @@ -3,7 +3,7 @@ import path from "path"; import fs from "fs/promises"; import EventEmitter from "events"; import { createWriteStream } from "fs"; -import { FindCursor } from "mongodb"; +import { AggregationCursor, FindCursor } from "mongodb"; import { EJSON, EJSONOptions, ObjectId } from "bson"; import { Transform } from "stream"; import { pipeline } from "stream/promises"; @@ -154,7 +154,7 @@ export class ExportsManager extends EventEmitter { exportTitle, jsonExportFormat, }: { - input: FindCursor; + input: FindCursor | AggregationCursor; exportName: string; exportTitle: string; jsonExportFormat: JSONExportFormat; @@ -194,7 +194,7 @@ export class ExportsManager extends EventEmitter { jsonExportFormat, inProgressExport, }: { - input: FindCursor; + input: FindCursor | AggregationCursor; jsonExportFormat: JSONExportFormat; inProgressExport: InProgressExport; }): Promise { diff --git a/src/tools/mongodb/read/export.ts b/src/tools/mongodb/read/export.ts index 9eaacba2f..3f85b41a6 100644 --- a/src/tools/mongodb/read/export.ts +++ b/src/tools/mongodb/read/export.ts @@ -1,19 +1,33 @@ import z from "zod"; import { ObjectId } from "bson"; +import { AggregationCursor, FindCursor } from "mongodb"; import { CallToolResult } from "@modelcontextprotocol/sdk/types.js"; import { OperationType, ToolArgs } from "../../tool.js"; import { DbOperationArgs, MongoDBToolBase } from "../mongodbTool.js"; import { FindArgs } from "./find.js"; import { jsonExportFormat } from "../../../common/exportsManager.js"; +import { AggregateArgs } from "./aggregate.js"; export class ExportTool extends MongoDBToolBase { public name = "export"; protected description = "Export a collection data or query results in the specified EJSON format."; protected argsShape = { - exportTitle: z.string().describe("A short description to uniquely identify the export."), ...DbOperationArgs, - ...FindArgs, - limit: z.number().optional().describe("The maximum number of documents to return"), + exportTitle: z.string().describe("A short description to uniquely identify the export."), + exportTarget: z + .array( + z.discriminatedUnion("name", [ + z.object({ + name: z.literal("find"), + arguments: z.object(FindArgs), + }), + z.object({ + name: z.literal("aggregate"), + arguments: z.object(AggregateArgs), + }), + ]) + ) + .describe("The export target along with its arguments."), jsonExportFormat: jsonExportFormat .default("relaxed") .describe( @@ -30,24 +44,37 @@ export class ExportTool extends MongoDBToolBase { database, collection, jsonExportFormat, - filter, - projection, - sort, - limit, exportTitle, + exportTarget: target, }: ToolArgs): Promise { const provider = await this.ensureConnected(); - const findCursor = provider.find(database, collection, filter ?? {}, { - projection, - sort, - limit, - promoteValues: false, - bsonRegExp: true, - }); + const exportTarget = target[0]; + if (!exportTarget) { + throw new Error("Export target not provided. Expected one of the following: `aggregate`, `find`"); + } + + let cursor: FindCursor | AggregationCursor; + if (exportTarget.name === "find") { + const { filter, projection, sort, limit } = exportTarget.arguments; + cursor = provider.find(database, collection, filter ?? {}, { + projection, + sort, + limit, + promoteValues: false, + bsonRegExp: true, + }); + } else { + const { pipeline } = exportTarget.arguments; + cursor = provider.aggregate(database, collection, pipeline, { + promoteValues: false, + bsonRegExp: true, + }); + } + const exportName = `${database}.${collection}.${new ObjectId().toString()}.json`; const { exportURI, exportPath } = await this.session.exportsManager.createJSONExport({ - input: findCursor, + input: cursor, exportName, exportTitle: exportTitle || diff --git a/tests/integration/resources/exportedData.test.ts b/tests/integration/resources/exportedData.test.ts index 94710d87f..3112fe878 100644 --- a/tests/integration/resources/exportedData.test.ts +++ b/tests/integration/resources/exportedData.test.ts @@ -65,7 +65,12 @@ describeWithMongoDB( await integration.connectMcpClient(); const exportResponse = await integration.mcpClient().callTool({ name: "export", - arguments: { database: "db", collection: "coll", exportTitle: "Export for db.coll" }, + arguments: { + database: "db", + collection: "coll", + exportTitle: "Export for db.coll", + exportTarget: [{ name: "find", arguments: {} }], + }, }); const exportedResourceURI = (exportResponse as CallToolResult).content.find( @@ -99,7 +104,12 @@ describeWithMongoDB( await integration.connectMcpClient(); const exportResponse = await integration.mcpClient().callTool({ name: "export", - arguments: { database: "db", collection: "coll", exportTitle: "Export for db.coll" }, + arguments: { + database: "db", + collection: "coll", + exportTitle: "Export for db.coll", + exportTarget: [{ name: "find", arguments: {} }], + }, }); const content = exportResponse.content as CallToolResult["content"]; const exportURI = contentWithResourceURILink(content)?.uri as string; @@ -122,7 +132,12 @@ describeWithMongoDB( await integration.connectMcpClient(); const exportResponse = await integration.mcpClient().callTool({ name: "export", - arguments: { database: "big", collection: "coll", exportTitle: "Export for big.coll" }, + arguments: { + database: "big", + collection: "coll", + exportTitle: "Export for big.coll", + exportTarget: [{ name: "find", arguments: {} }], + }, }); const content = exportResponse.content as CallToolResult["content"]; const exportURI = contentWithResourceURILink(content)?.uri as string; diff --git a/tests/integration/tools/mongodb/read/export.test.ts b/tests/integration/tools/mongodb/read/export.test.ts index 343f3ef45..f02460f01 100644 --- a/tests/integration/tools/mongodb/read/export.test.ts +++ b/tests/integration/tools/mongodb/read/export.test.ts @@ -62,12 +62,6 @@ describeWithMongoDB( type: "string", required: true, }, - { - name: "filter", - description: "The query filter, matching the syntax of the query argument of db.collection.find()", - type: "object", - required: false, - }, { name: "jsonExportFormat", description: [ @@ -79,24 +73,10 @@ describeWithMongoDB( required: false, }, { - name: "limit", - description: "The maximum number of documents to return", - type: "number", - required: false, - }, - { - name: "projection", - description: - "The projection, matching the syntax of the projection argument of db.collection.find()", - type: "object", - required: false, - }, - { - name: "sort", - description: - "A document, describing the sort order, matching the syntax of the sort argument of cursor.sort(). The keys of the object are the fields to sort on, while the values are the sort directions (1 for ascending, -1 for descending).", - type: "object", - required: false, + name: "exportTarget", + type: "array", + description: "The export target along with its arguments.", + required: true, }, ] ); @@ -126,6 +106,14 @@ describeWithMongoDB( database: "non-existent", collection: "foos", exportTitle: "Export for non-existent.foos", + exportTarget: [ + { + name: "find", + arguments: { + filter: {}, + }, + }, + ], }, }); const content = response.content as CallToolResult["content"]; @@ -165,6 +153,14 @@ describeWithMongoDB( database: integration.randomDbName(), collection: "foo", exportTitle: `Export for ${integration.randomDbName()}.foo`, + exportTarget: [ + { + name: "find", + arguments: { + filter: {}, + }, + }, + ], }, }); const content = response.content as CallToolResult["content"]; @@ -192,8 +188,15 @@ describeWithMongoDB( arguments: { database: integration.randomDbName(), collection: "foo", - filter: { name: "foo" }, exportTitle: `Export for ${integration.randomDbName()}.foo`, + exportTarget: [ + { + name: "find", + arguments: { + filter: { name: "foo" }, + }, + }, + ], }, }); const content = response.content as CallToolResult["content"]; @@ -220,8 +223,16 @@ describeWithMongoDB( arguments: { database: integration.randomDbName(), collection: "foo", - limit: 1, exportTitle: `Export for ${integration.randomDbName()}.foo`, + exportTarget: [ + { + name: "find", + arguments: { + filter: {}, + limit: 1, + }, + }, + ], }, }); const content = response.content as CallToolResult["content"]; @@ -248,9 +259,17 @@ describeWithMongoDB( arguments: { database: integration.randomDbName(), collection: "foo", - limit: 1, - sort: { longNumber: 1 }, exportTitle: `Export for ${integration.randomDbName()}.foo`, + exportTarget: [ + { + name: "find", + arguments: { + filter: {}, + limit: 1, + sort: { longNumber: 1 }, + }, + }, + ], }, }); const content = response.content as CallToolResult["content"]; @@ -277,9 +296,17 @@ describeWithMongoDB( arguments: { database: integration.randomDbName(), collection: "foo", - limit: 1, - projection: { _id: 0, name: 1 }, exportTitle: `Export for ${integration.randomDbName()}.foo`, + exportTarget: [ + { + name: "find", + arguments: { + filter: {}, + limit: 1, + projection: { _id: 0, name: 1 }, + }, + }, + ], }, }); const content = response.content as CallToolResult["content"]; @@ -309,10 +336,18 @@ describeWithMongoDB( arguments: { database: integration.randomDbName(), collection: "foo", - limit: 1, - projection: { _id: 0 }, jsonExportFormat: "relaxed", exportTitle: `Export for ${integration.randomDbName()}.foo`, + exportTarget: [ + { + name: "find", + arguments: { + filter: {}, + limit: 1, + projection: { _id: 0 }, + }, + }, + ], }, }); const content = response.content as CallToolResult["content"]; @@ -343,10 +378,18 @@ describeWithMongoDB( arguments: { database: integration.randomDbName(), collection: "foo", - limit: 1, - projection: { _id: 0 }, jsonExportFormat: "canonical", exportTitle: `Export for ${integration.randomDbName()}.foo`, + exportTarget: [ + { + name: "find", + arguments: { + filter: {}, + limit: 1, + projection: { _id: 0 }, + }, + }, + ], }, }); const content = response.content as CallToolResult["content"]; @@ -371,6 +414,48 @@ describeWithMongoDB( }, ]); }); + + it("should allow exporting an aggregation", async () => { + await integration.connectMcpClient(); + const response = await integration.mcpClient().callTool({ + name: "export", + arguments: { + database: integration.randomDbName(), + collection: "foo", + exportTitle: `Export for ${integration.randomDbName()}.foo`, + exportTarget: [ + { + name: "aggregate", + arguments: { + pipeline: [ + { + $match: {}, + }, + { + $limit: 1, + }, + ], + }, + }, + ], + }, + }); + const content = response.content as CallToolResult["content"]; + const exportURI = contentWithResourceURILink(content)?.uri as string; + await resourceChangedNotification(integration.mcpClient(), exportURI); + + const localPathPart = contentWithExportPath(content); + expect(localPathPart).toBeDefined(); + const [, localPath] = /"(.*)"/.exec(String(localPathPart?.text)) ?? []; + expect(localPath).toBeDefined(); + + const exportedContent = JSON.parse(await fs.readFile(localPath as string, "utf8")) as Record< + string, + unknown + >[]; + expect(exportedContent).toHaveLength(1); + expect(exportedContent[0]?.name).toEqual("foo"); + }); }); }, () => userConfig From df9927a71e1526b4accccac4190739ed7a3357d8 Mon Sep 17 00:00:00 2001 From: Himanshu Singh Date: Wed, 20 Aug 2025 10:55:37 +0200 Subject: [PATCH 2/6] chore: update accuracy tests --- tests/accuracy/export.test.ts | 61 ++++++++++++++++++++++++++--------- 1 file changed, 45 insertions(+), 16 deletions(-) diff --git a/tests/accuracy/export.test.ts b/tests/accuracy/export.test.ts index 9e1f0cff5..a4f7550e5 100644 --- a/tests/accuracy/export.test.ts +++ b/tests/accuracy/export.test.ts @@ -10,8 +10,13 @@ describeAccuracyTests([ parameters: { database: "mflix", collection: "movies", - filter: Matcher.emptyObjectOrUndefined, - limit: Matcher.undefined, + exportTitle: Matcher.string(), + exportTarget: [ + { + name: "find", + arguments: {}, + }, + ], }, }, ], @@ -24,9 +29,17 @@ describeAccuracyTests([ parameters: { database: "mflix", collection: "movies", - filter: { - runtime: { $lt: 100 }, - }, + exportTitle: Matcher.string(), + exportTarget: [ + { + name: "find", + arguments: { + filter: { + runtime: { $lt: 100 }, + }, + }, + }, + ], }, }, ], @@ -39,14 +52,22 @@ describeAccuracyTests([ parameters: { database: "mflix", collection: "movies", - projection: { - title: 1, - _id: Matcher.anyOf( - Matcher.undefined, - Matcher.number((value) => value === 0) - ), - }, - filter: Matcher.emptyObjectOrUndefined, + exportTitle: Matcher.string(), + exportTarget: [ + { + name: "find", + arguments: { + projection: { + title: 1, + _id: Matcher.anyOf( + Matcher.undefined, + Matcher.number((value) => value === 0) + ), + }, + filter: Matcher.emptyObjectOrUndefined, + }, + }, + ], }, }, ], @@ -59,9 +80,17 @@ describeAccuracyTests([ parameters: { database: "mflix", collection: "movies", - filter: { genres: "Horror" }, - sort: { runtime: 1 }, - limit: 2, + exportTitle: Matcher.string(), + exportTarget: [ + { + name: "find", + arguments: { + filter: { genres: "Horror" }, + sort: { runtime: 1 }, + limit: 2, + }, + }, + ], }, }, ], From 04c4e9fd0786c534aafba24158c8b1f41b2ae33a Mon Sep 17 00:00:00 2001 From: Himanshu Singh Date: Wed, 20 Aug 2025 14:24:53 +0200 Subject: [PATCH 3/6] chore: remove default from limit arg in find target --- src/tools/mongodb/read/export.ts | 7 +++++-- src/tools/mongodb/read/find.ts | 4 +++- 2 files changed, 8 insertions(+), 3 deletions(-) diff --git a/src/tools/mongodb/read/export.ts b/src/tools/mongodb/read/export.ts index 3f85b41a6..8b044272f 100644 --- a/src/tools/mongodb/read/export.ts +++ b/src/tools/mongodb/read/export.ts @@ -4,7 +4,7 @@ import { AggregationCursor, FindCursor } from "mongodb"; import { CallToolResult } from "@modelcontextprotocol/sdk/types.js"; import { OperationType, ToolArgs } from "../../tool.js"; import { DbOperationArgs, MongoDBToolBase } from "../mongodbTool.js"; -import { FindArgs } from "./find.js"; +import { FindArgs, limitArg } from "./find.js"; import { jsonExportFormat } from "../../../common/exportsManager.js"; import { AggregateArgs } from "./aggregate.js"; @@ -19,7 +19,10 @@ export class ExportTool extends MongoDBToolBase { z.discriminatedUnion("name", [ z.object({ name: z.literal("find"), - arguments: z.object(FindArgs), + arguments: z.object({ + ...FindArgs, + limit: limitArg, + }), }), z.object({ name: z.literal("aggregate"), diff --git a/src/tools/mongodb/read/find.ts b/src/tools/mongodb/read/find.ts index f04c87f6c..648878852 100644 --- a/src/tools/mongodb/read/find.ts +++ b/src/tools/mongodb/read/find.ts @@ -5,6 +5,8 @@ import { ToolArgs, OperationType } from "../../tool.js"; import { SortDirection } from "mongodb"; import { checkIndexUsage } from "../../../helpers/indexCheck.js"; +export const limitArg = z.number().optional().describe("The maximum number of documents to return"); + export const FindArgs = { filter: z .object({}) @@ -16,7 +18,7 @@ export const FindArgs = { .passthrough() .optional() .describe("The projection, matching the syntax of the projection argument of db.collection.find()"), - limit: z.number().optional().default(10).describe("The maximum number of documents to return"), + limit: limitArg.default(10), sort: z .object({}) .catchall(z.custom()) From b226827b7b3444ce1c17b6b4245596e988043dee Mon Sep 17 00:00:00 2001 From: Himanshu Singh Date: Wed, 20 Aug 2025 20:23:06 +0200 Subject: [PATCH 4/6] chore: use Zod.removeDefault() for removing default from limit arg --- src/tools/mongodb/read/export.ts | 4 ++-- src/tools/mongodb/read/find.ts | 4 +--- 2 files changed, 3 insertions(+), 5 deletions(-) diff --git a/src/tools/mongodb/read/export.ts b/src/tools/mongodb/read/export.ts index 8b044272f..a2ffa2fab 100644 --- a/src/tools/mongodb/read/export.ts +++ b/src/tools/mongodb/read/export.ts @@ -4,7 +4,7 @@ import { AggregationCursor, FindCursor } from "mongodb"; import { CallToolResult } from "@modelcontextprotocol/sdk/types.js"; import { OperationType, ToolArgs } from "../../tool.js"; import { DbOperationArgs, MongoDBToolBase } from "../mongodbTool.js"; -import { FindArgs, limitArg } from "./find.js"; +import { FindArgs } from "./find.js"; import { jsonExportFormat } from "../../../common/exportsManager.js"; import { AggregateArgs } from "./aggregate.js"; @@ -21,7 +21,7 @@ export class ExportTool extends MongoDBToolBase { name: z.literal("find"), arguments: z.object({ ...FindArgs, - limit: limitArg, + limit: FindArgs.limit.removeDefault(), }), }), z.object({ diff --git a/src/tools/mongodb/read/find.ts b/src/tools/mongodb/read/find.ts index 648878852..f04c87f6c 100644 --- a/src/tools/mongodb/read/find.ts +++ b/src/tools/mongodb/read/find.ts @@ -5,8 +5,6 @@ import { ToolArgs, OperationType } from "../../tool.js"; import { SortDirection } from "mongodb"; import { checkIndexUsage } from "../../../helpers/indexCheck.js"; -export const limitArg = z.number().optional().describe("The maximum number of documents to return"); - export const FindArgs = { filter: z .object({}) @@ -18,7 +16,7 @@ export const FindArgs = { .passthrough() .optional() .describe("The projection, matching the syntax of the projection argument of db.collection.find()"), - limit: limitArg.default(10), + limit: z.number().optional().default(10).describe("The maximum number of documents to return"), sort: z .object({}) .catchall(z.custom()) From 5c2638f1c45e3956961cccb494d2870a78e45b5e Mon Sep 17 00:00:00 2001 From: Himanshu Singh Date: Wed, 20 Aug 2025 20:23:27 +0200 Subject: [PATCH 5/6] chore: allow aggregations to use disk --- src/tools/mongodb/read/export.ts | 1 + 1 file changed, 1 insertion(+) diff --git a/src/tools/mongodb/read/export.ts b/src/tools/mongodb/read/export.ts index a2ffa2fab..1777f74e3 100644 --- a/src/tools/mongodb/read/export.ts +++ b/src/tools/mongodb/read/export.ts @@ -71,6 +71,7 @@ export class ExportTool extends MongoDBToolBase { cursor = provider.aggregate(database, collection, pipeline, { promoteValues: false, bsonRegExp: true, + allowDiskUse: true, }); } From 509c37b86d26e39353b2c30abd74f70912050ead Mon Sep 17 00:00:00 2001 From: Himanshu Singh Date: Wed, 20 Aug 2025 21:01:38 +0200 Subject: [PATCH 6/6] chore: adds accuracy test for aggregation export --- src/tools/mongodb/read/export.ts | 20 +++++++++++++------- tests/accuracy/export.test.ts | 30 ++++++++++++++++++++++++++++++ 2 files changed, 43 insertions(+), 7 deletions(-) diff --git a/src/tools/mongodb/read/export.ts b/src/tools/mongodb/read/export.ts index 1777f74e3..2a6097c8e 100644 --- a/src/tools/mongodb/read/export.ts +++ b/src/tools/mongodb/read/export.ts @@ -18,15 +18,21 @@ export class ExportTool extends MongoDBToolBase { .array( z.discriminatedUnion("name", [ z.object({ - name: z.literal("find"), - arguments: z.object({ - ...FindArgs, - limit: FindArgs.limit.removeDefault(), - }), + name: z + .literal("find") + .describe("The literal name 'find' to represent a find cursor as target."), + arguments: z + .object({ + ...FindArgs, + limit: FindArgs.limit.removeDefault(), + }) + .describe("The arguments for 'find' operation."), }), z.object({ - name: z.literal("aggregate"), - arguments: z.object(AggregateArgs), + name: z + .literal("aggregate") + .describe("The literal name 'aggregate' to represent an aggregation cursor as target."), + arguments: z.object(AggregateArgs).describe("The arguments for 'aggregate' operation."), }), ]) ) diff --git a/tests/accuracy/export.test.ts b/tests/accuracy/export.test.ts index a4f7550e5..5b2624171 100644 --- a/tests/accuracy/export.test.ts +++ b/tests/accuracy/export.test.ts @@ -95,4 +95,34 @@ describeAccuracyTests([ }, ], }, + { + prompt: "Export an aggregation that groups all movie titles by the field release_year from mflix.movies", + expectedToolCalls: [ + { + toolName: "export", + parameters: { + database: "mflix", + collection: "movies", + exportTitle: Matcher.string(), + exportTarget: [ + { + name: "aggregate", + arguments: { + pipeline: [ + { + $group: { + _id: "$release_year", + titles: { + $push: "$title", + }, + }, + }, + ], + }, + }, + ], + }, + }, + ], + }, ]);