Skip to content

Navigation Menu

Sign in
Appearance settings

Search code, repositories, users, issues, pull requests...

Provide feedback

We read every piece of feedback, and take your input very seriously.

Saved searches

Use saved searches to filter your results more quickly
Sign up
Appearance settings

Commit c6f4b01

Browse files
feat: adds support for exporting aggregations MCP-16 (#464)
1 parent c4ba2c9 commit c6f4b01

File tree

5 files changed

+267
-71
lines changed

5 files changed

+267
-71
lines changed

‎src/common/exportsManager.ts

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@ import path from "path";
33
import fs from "fs/promises";
44
import EventEmitter from "events";
55
import { createWriteStream } from "fs";
6-
import { FindCursor } from "mongodb";
6+
import { AggregationCursor,FindCursor } from "mongodb";
77
import { EJSON, EJSONOptions, ObjectId } from "bson";
88
import { Transform } from "stream";
99
import { pipeline } from "stream/promises";
@@ -154,7 +154,7 @@ export class ExportsManager extends EventEmitter<ExportsManagerEvents> {
154154
exportTitle,
155155
jsonExportFormat,
156156
}: {
157-
input: FindCursor;
157+
input: FindCursor|AggregationCursor;
158158
exportName: string;
159159
exportTitle: string;
160160
jsonExportFormat: JSONExportFormat;
@@ -194,7 +194,7 @@ export class ExportsManager extends EventEmitter<ExportsManagerEvents> {
194194
jsonExportFormat,
195195
inProgressExport,
196196
}: {
197-
input: FindCursor;
197+
input: FindCursor|AggregationCursor;
198198
jsonExportFormat: JSONExportFormat;
199199
inProgressExport: InProgressExport;
200200
}): Promise<void> {

‎src/tools/mongodb/read/export.ts

Lines changed: 52 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -1,19 +1,42 @@
11
import z from "zod";
22
import { ObjectId } from "bson";
3+
import { AggregationCursor, FindCursor } from "mongodb";
34
import { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
45
import { OperationType, ToolArgs } from "../../tool.js";
56
import { DbOperationArgs, MongoDBToolBase } from "../mongodbTool.js";
67
import { FindArgs } from "./find.js";
78
import { jsonExportFormat } from "../../../common/exportsManager.js";
9+
import { AggregateArgs } from "./aggregate.js";
810

911
export class ExportTool extends MongoDBToolBase {
1012
public name = "export";
1113
protected description = "Export a collection data or query results in the specified EJSON format.";
1214
protected argsShape = {
13-
exportTitle: z.string().describe("A short description to uniquely identify the export."),
1415
...DbOperationArgs,
15-
...FindArgs,
16-
limit: z.number().optional().describe("The maximum number of documents to return"),
16+
exportTitle: z.string().describe("A short description to uniquely identify the export."),
17+
exportTarget: z
18+
.array(
19+
z.discriminatedUnion("name", [
20+
z.object({
21+
name: z
22+
.literal("find")
23+
.describe("The literal name 'find' to represent a find cursor as target."),
24+
arguments: z
25+
.object({
26+
...FindArgs,
27+
limit: FindArgs.limit.removeDefault(),
28+
})
29+
.describe("The arguments for 'find' operation."),
30+
}),
31+
z.object({
32+
name: z
33+
.literal("aggregate")
34+
.describe("The literal name 'aggregate' to represent an aggregation cursor as target."),
35+
arguments: z.object(AggregateArgs).describe("The arguments for 'aggregate' operation."),
36+
}),
37+
])
38+
)
39+
.describe("The export target along with its arguments."),
1740
jsonExportFormat: jsonExportFormat
1841
.default("relaxed")
1942
.describe(
@@ -30,24 +53,38 @@ export class ExportTool extends MongoDBToolBase {
3053
database,
3154
collection,
3255
jsonExportFormat,
33-
filter,
34-
projection,
35-
sort,
36-
limit,
3756
exportTitle,
57+
exportTarget: target,
3858
}: ToolArgs<typeof this.argsShape>): Promise<CallToolResult> {
3959
const provider = await this.ensureConnected();
40-
const findCursor = provider.find(database, collection, filter ?? {}, {
41-
projection,
42-
sort,
43-
limit,
44-
promoteValues: false,
45-
bsonRegExp: true,
46-
});
60+
const exportTarget = target[0];
61+
if (!exportTarget) {
62+
throw new Error("Export target not provided. Expected one of the following: `aggregate`, `find`");
63+
}
64+
65+
let cursor: FindCursor | AggregationCursor;
66+
if (exportTarget.name === "find") {
67+
const { filter, projection, sort, limit } = exportTarget.arguments;
68+
cursor = provider.find(database, collection, filter ?? {}, {
69+
projection,
70+
sort,
71+
limit,
72+
promoteValues: false,
73+
bsonRegExp: true,
74+
});
75+
} else {
76+
const { pipeline } = exportTarget.arguments;
77+
cursor = provider.aggregate(database, collection, pipeline, {
78+
promoteValues: false,
79+
bsonRegExp: true,
80+
allowDiskUse: true,
81+
});
82+
}
83+
4784
const exportName = `${database}.${collection}.${new ObjectId().toString()}.json`;
4885

4986
const { exportURI, exportPath } = await this.session.exportsManager.createJSONExport({
50-
input: findCursor,
87+
input: cursor,
5188
exportName,
5289
exportTitle:
5390
exportTitle ||

‎tests/accuracy/export.test.ts

Lines changed: 75 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -10,8 +10,13 @@ describeAccuracyTests([
1010
parameters: {
1111
database: "mflix",
1212
collection: "movies",
13-
filter: Matcher.emptyObjectOrUndefined,
14-
limit: Matcher.undefined,
13+
exportTitle: Matcher.string(),
14+
exportTarget: [
15+
{
16+
name: "find",
17+
arguments: {},
18+
},
19+
],
1520
},
1621
},
1722
],
@@ -24,9 +29,17 @@ describeAccuracyTests([
2429
parameters: {
2530
database: "mflix",
2631
collection: "movies",
27-
filter: {
28-
runtime: { $lt: 100 },
29-
},
32+
exportTitle: Matcher.string(),
33+
exportTarget: [
34+
{
35+
name: "find",
36+
arguments: {
37+
filter: {
38+
runtime: { $lt: 100 },
39+
},
40+
},
41+
},
42+
],
3043
},
3144
},
3245
],
@@ -39,14 +52,22 @@ describeAccuracyTests([
3952
parameters: {
4053
database: "mflix",
4154
collection: "movies",
42-
projection: {
43-
title: 1,
44-
_id: Matcher.anyOf(
45-
Matcher.undefined,
46-
Matcher.number((value) => value === 0)
47-
),
48-
},
49-
filter: Matcher.emptyObjectOrUndefined,
55+
exportTitle: Matcher.string(),
56+
exportTarget: [
57+
{
58+
name: "find",
59+
arguments: {
60+
projection: {
61+
title: 1,
62+
_id: Matcher.anyOf(
63+
Matcher.undefined,
64+
Matcher.number((value) => value === 0)
65+
),
66+
},
67+
filter: Matcher.emptyObjectOrUndefined,
68+
},
69+
},
70+
],
5071
},
5172
},
5273
],
@@ -59,9 +80,47 @@ describeAccuracyTests([
5980
parameters: {
6081
database: "mflix",
6182
collection: "movies",
62-
filter: { genres: "Horror" },
63-
sort: { runtime: 1 },
64-
limit: 2,
83+
exportTitle: Matcher.string(),
84+
exportTarget: [
85+
{
86+
name: "find",
87+
arguments: {
88+
filter: { genres: "Horror" },
89+
sort: { runtime: 1 },
90+
limit: 2,
91+
},
92+
},
93+
],
94+
},
95+
},
96+
],
97+
},
98+
{
99+
prompt: "Export an aggregation that groups all movie titles by the field release_year from mflix.movies",
100+
expectedToolCalls: [
101+
{
102+
toolName: "export",
103+
parameters: {
104+
database: "mflix",
105+
collection: "movies",
106+
exportTitle: Matcher.string(),
107+
exportTarget: [
108+
{
109+
name: "aggregate",
110+
arguments: {
111+
pipeline: [
112+
{
113+
$group: {
114+
_id: "$release_year",
115+
titles: {
116+
$push: "$title",
117+
},
118+
},
119+
},
120+
],
121+
},
122+
},
123+
],
65124
},
66125
},
67126
],

‎tests/integration/resources/exportedData.test.ts

Lines changed: 18 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -65,7 +65,12 @@ describeWithMongoDB(
6565
await integration.connectMcpClient();
6666
const exportResponse = await integration.mcpClient().callTool({
6767
name: "export",
68-
arguments: { database: "db", collection: "coll", exportTitle: "Export for db.coll" },
68+
arguments: {
69+
database: "db",
70+
collection: "coll",
71+
exportTitle: "Export for db.coll",
72+
exportTarget: [{ name: "find", arguments: {} }],
73+
},
6974
});
7075

7176
const exportedResourceURI = (exportResponse as CallToolResult).content.find(
@@ -99,7 +104,12 @@ describeWithMongoDB(
99104
await integration.connectMcpClient();
100105
const exportResponse = await integration.mcpClient().callTool({
101106
name: "export",
102-
arguments: { database: "db", collection: "coll", exportTitle: "Export for db.coll" },
107+
arguments: {
108+
database: "db",
109+
collection: "coll",
110+
exportTitle: "Export for db.coll",
111+
exportTarget: [{ name: "find", arguments: {} }],
112+
},
103113
});
104114
const content = exportResponse.content as CallToolResult["content"];
105115
const exportURI = contentWithResourceURILink(content)?.uri as string;
@@ -122,7 +132,12 @@ describeWithMongoDB(
122132
await integration.connectMcpClient();
123133
const exportResponse = await integration.mcpClient().callTool({
124134
name: "export",
125-
arguments: { database: "big", collection: "coll", exportTitle: "Export for big.coll" },
135+
arguments: {
136+
database: "big",
137+
collection: "coll",
138+
exportTitle: "Export for big.coll",
139+
exportTarget: [{ name: "find", arguments: {} }],
140+
},
126141
});
127142
const content = exportResponse.content as CallToolResult["content"];
128143
const exportURI = contentWithResourceURILink(content)?.uri as string;

0 commit comments

Comments
(0)

AltStyle によって変換されたページ (->オリジナル) /