Skip to content

feat: adds support for exporting aggregations MCP-16 #464

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.

Sign up for GitHub

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

Merged
merged 6 commits into from
Aug 21, 2025
Merged

feat: adds support for exporting aggregations MCP-16 #464

Show file tree
Hide file tree
Changes from 3 commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
293cfe2
feat: adds support for exporting aggregations
himanshusinghs Aug 20, 2025
df9927a
chore: update accuracy tests
himanshusinghs Aug 20, 2025
04c4e9f
chore: remove default from limit arg in find target
himanshusinghs Aug 20, 2025
b226827
chore: use Zod.removeDefault() for removing default from limit arg
himanshusinghs Aug 20, 2025
5c2638f
chore: allow aggregations to use disk
himanshusinghs Aug 20, 2025
509c37b
chore: adds accuracy test for aggregation export
himanshusinghs Aug 20, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 3 additions & 3 deletions src/common/exportsManager.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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";
Expand Down Expand Up @@ -154,7 +154,7 @@ export class ExportsManager extends EventEmitter<ExportsManagerEvents> {
exportTitle,
jsonExportFormat,
}: {
input: FindCursor;
input: FindCursor | AggregationCursor;
exportName: string;
exportTitle: string;
jsonExportFormat: JSONExportFormat;
Expand Down Expand Up @@ -194,7 +194,7 @@ export class ExportsManager extends EventEmitter<ExportsManagerEvents> {
jsonExportFormat,
inProgressExport,
}: {
input: FindCursor;
input: FindCursor | AggregationCursor;
jsonExportFormat: JSONExportFormat;
inProgressExport: InProgressExport;
}): Promise<void> {
Expand Down
62 changes: 46 additions & 16 deletions src/tools/mongodb/read/export.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,19 +1,36 @@
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 { FindArgs, limitArg } 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(
Comment on lines +17 to +18
Copy link
Collaborator

Choose a reason for hiding this comment

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

I wonder if there's a way to avoid using an array here - I remember that was a workaround, but perhaps a newer version of the server/zod has addressed that.

Copy link
Collaborator Author

Choose a reason for hiding this comment

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

We can already use Zod's discriminated unions without wrapping them in array. I thought this was done to avoid having to expand the validateMetadata tests back then because that test function is currently strictly bound to work with normal parameters. I also brushed off that effort only because it was not that important at the moment and instead created a ticket MCP-128 to follow up on that.

z.discriminatedUnion("name", [
z.object({
name: z.literal("find"),
arguments: z.object({
...FindArgs,
limit: limitArg,
}),
}),
z.object({
name: z.literal("aggregate"),
arguments: z.object(AggregateArgs),
}),
])
)
.describe("The export target along with its arguments."),
jsonExportFormat: jsonExportFormat
.default("relaxed")
.describe(
Expand All @@ -30,24 +47,37 @@ export class ExportTool extends MongoDBToolBase {
database,
collection,
jsonExportFormat,
filter,
projection,
sort,
limit,
exportTitle,
exportTarget: target,
}: ToolArgs<typeof this.argsShape>): Promise<CallToolResult> {
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 ||
Expand Down
4 changes: 3 additions & 1 deletion src/tools/mongodb/read/find.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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({})
Expand All @@ -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<SortDirection>())
Expand Down
61 changes: 45 additions & 16 deletions tests/accuracy/export.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,8 +10,13 @@ describeAccuracyTests([
parameters: {
database: "mflix",
collection: "movies",
filter: Matcher.emptyObjectOrUndefined,
limit: Matcher.undefined,
exportTitle: Matcher.string(),
exportTarget: [
{
name: "find",
arguments: {},
},
],
},
},
],
Expand All @@ -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 },
},
},
},
],
},
},
],
Expand All @@ -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,
},
},
],
},
},
],
Expand All @@ -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,
},
},
],
},
},
],
Expand Down
21 changes: 18 additions & 3 deletions tests/integration/resources/exportedData.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down Expand Up @@ -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;
Expand All @@ -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;
Expand Down
Loading
Loading