Skip to content

Commit 293cfe2

Browse files
himanshusinghshimanshusinghs
committed
feat: adds support for exporting aggregations
1 parent c4ba2c9 commit 293cfe2

File tree

4 files changed

+182
-55
lines changed

4 files changed

+182
-55
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: 42 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -1,19 +1,33 @@
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.literal("find"),
22+
arguments: z.object(FindArgs),
23+
}),
24+
z.object({
25+
name: z.literal("aggregate"),
26+
arguments: z.object(AggregateArgs),
27+
}),
28+
])
29+
)
30+
.describe("The export target along with its arguments."),
1731
jsonExportFormat: jsonExportFormat
1832
.default("relaxed")
1933
.describe(
@@ -30,24 +44,37 @@ export class ExportTool extends MongoDBToolBase {
3044
database,
3145
collection,
3246
jsonExportFormat,
33-
filter,
34-
projection,
35-
sort,
36-
limit,
3747
exportTitle,
48+
exportTarget: target,
3849
}: ToolArgs<typeof this.argsShape>): Promise<CallToolResult> {
3950
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-
});
51+
const exportTarget = target[0];
52+
if (!exportTarget) {
53+
throw new Error("Export target not provided. Expected one of the following: `aggregate`, `find`");
54+
}
55+
56+
let cursor: FindCursor | AggregationCursor;
57+
if (exportTarget.name === "find") {
58+
const { filter, projection, sort, limit } = exportTarget.arguments;
59+
cursor = provider.find(database, collection, filter ?? {}, {
60+
projection,
61+
sort,
62+
limit,
63+
promoteValues: false,
64+
bsonRegExp: true,
65+
});
66+
} else {
67+
const { pipeline } = exportTarget.arguments;
68+
cursor = provider.aggregate(database, collection, pipeline, {
69+
promoteValues: false,
70+
bsonRegExp: true,
71+
});
72+
}
73+
4774
const exportName = `${database}.${collection}.${new ObjectId().toString()}.json`;
4875

4976
const { exportURI, exportPath } = await this.session.exportsManager.createJSONExport({
50-
input: findCursor,
77+
input: cursor,
5178
exportName,
5279
exportTitle:
5380
exportTitle ||

‎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;

‎tests/integration/tools/mongodb/read/export.test.ts‎

Lines changed: 119 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -62,12 +62,6 @@ describeWithMongoDB(
6262
type: "string",
6363
required: true,
6464
},
65-
{
66-
name: "filter",
67-
description: "The query filter, matching the syntax of the query argument of db.collection.find()",
68-
type: "object",
69-
required: false,
70-
},
7165
{
7266
name: "jsonExportFormat",
7367
description: [
@@ -79,24 +73,10 @@ describeWithMongoDB(
7973
required: false,
8074
},
8175
{
82-
name: "limit",
83-
description: "The maximum number of documents to return",
84-
type: "number",
85-
required: false,
86-
},
87-
{
88-
name: "projection",
89-
description:
90-
"The projection, matching the syntax of the projection argument of db.collection.find()",
91-
type: "object",
92-
required: false,
93-
},
94-
{
95-
name: "sort",
96-
description:
97-
"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).",
98-
type: "object",
99-
required: false,
76+
name: "exportTarget",
77+
type: "array",
78+
description: "The export target along with its arguments.",
79+
required: true,
10080
},
10181
]
10282
);
@@ -126,6 +106,14 @@ describeWithMongoDB(
126106
database: "non-existent",
127107
collection: "foos",
128108
exportTitle: "Export for non-existent.foos",
109+
exportTarget: [
110+
{
111+
name: "find",
112+
arguments: {
113+
filter: {},
114+
},
115+
},
116+
],
129117
},
130118
});
131119
const content = response.content as CallToolResult["content"];
@@ -165,6 +153,14 @@ describeWithMongoDB(
165153
database: integration.randomDbName(),
166154
collection: "foo",
167155
exportTitle: `Export for ${integration.randomDbName()}.foo`,
156+
exportTarget: [
157+
{
158+
name: "find",
159+
arguments: {
160+
filter: {},
161+
},
162+
},
163+
],
168164
},
169165
});
170166
const content = response.content as CallToolResult["content"];
@@ -192,8 +188,15 @@ describeWithMongoDB(
192188
arguments: {
193189
database: integration.randomDbName(),
194190
collection: "foo",
195-
filter: { name: "foo" },
196191
exportTitle: `Export for ${integration.randomDbName()}.foo`,
192+
exportTarget: [
193+
{
194+
name: "find",
195+
arguments: {
196+
filter: { name: "foo" },
197+
},
198+
},
199+
],
197200
},
198201
});
199202
const content = response.content as CallToolResult["content"];
@@ -220,8 +223,16 @@ describeWithMongoDB(
220223
arguments: {
221224
database: integration.randomDbName(),
222225
collection: "foo",
223-
limit: 1,
224226
exportTitle: `Export for ${integration.randomDbName()}.foo`,
227+
exportTarget: [
228+
{
229+
name: "find",
230+
arguments: {
231+
filter: {},
232+
limit: 1,
233+
},
234+
},
235+
],
225236
},
226237
});
227238
const content = response.content as CallToolResult["content"];
@@ -248,9 +259,17 @@ describeWithMongoDB(
248259
arguments: {
249260
database: integration.randomDbName(),
250261
collection: "foo",
251-
limit: 1,
252-
sort: { longNumber: 1 },
253262
exportTitle: `Export for ${integration.randomDbName()}.foo`,
263+
exportTarget: [
264+
{
265+
name: "find",
266+
arguments: {
267+
filter: {},
268+
limit: 1,
269+
sort: { longNumber: 1 },
270+
},
271+
},
272+
],
254273
},
255274
});
256275
const content = response.content as CallToolResult["content"];
@@ -277,9 +296,17 @@ describeWithMongoDB(
277296
arguments: {
278297
database: integration.randomDbName(),
279298
collection: "foo",
280-
limit: 1,
281-
projection: { _id: 0, name: 1 },
282299
exportTitle: `Export for ${integration.randomDbName()}.foo`,
300+
exportTarget: [
301+
{
302+
name: "find",
303+
arguments: {
304+
filter: {},
305+
limit: 1,
306+
projection: { _id: 0, name: 1 },
307+
},
308+
},
309+
],
283310
},
284311
});
285312
const content = response.content as CallToolResult["content"];
@@ -309,10 +336,18 @@ describeWithMongoDB(
309336
arguments: {
310337
database: integration.randomDbName(),
311338
collection: "foo",
312-
limit: 1,
313-
projection: { _id: 0 },
314339
jsonExportFormat: "relaxed",
315340
exportTitle: `Export for ${integration.randomDbName()}.foo`,
341+
exportTarget: [
342+
{
343+
name: "find",
344+
arguments: {
345+
filter: {},
346+
limit: 1,
347+
projection: { _id: 0 },
348+
},
349+
},
350+
],
316351
},
317352
});
318353
const content = response.content as CallToolResult["content"];
@@ -343,10 +378,18 @@ describeWithMongoDB(
343378
arguments: {
344379
database: integration.randomDbName(),
345380
collection: "foo",
346-
limit: 1,
347-
projection: { _id: 0 },
348381
jsonExportFormat: "canonical",
349382
exportTitle: `Export for ${integration.randomDbName()}.foo`,
383+
exportTarget: [
384+
{
385+
name: "find",
386+
arguments: {
387+
filter: {},
388+
limit: 1,
389+
projection: { _id: 0 },
390+
},
391+
},
392+
],
350393
},
351394
});
352395
const content = response.content as CallToolResult["content"];
@@ -371,6 +414,48 @@ describeWithMongoDB(
371414
},
372415
]);
373416
});
417+
418+
it("should allow exporting an aggregation", async () => {
419+
await integration.connectMcpClient();
420+
const response = await integration.mcpClient().callTool({
421+
name: "export",
422+
arguments: {
423+
database: integration.randomDbName(),
424+
collection: "foo",
425+
exportTitle: `Export for ${integration.randomDbName()}.foo`,
426+
exportTarget: [
427+
{
428+
name: "aggregate",
429+
arguments: {
430+
pipeline: [
431+
{
432+
$match: {},
433+
},
434+
{
435+
$limit: 1,
436+
},
437+
],
438+
},
439+
},
440+
],
441+
},
442+
});
443+
const content = response.content as CallToolResult["content"];
444+
const exportURI = contentWithResourceURILink(content)?.uri as string;
445+
await resourceChangedNotification(integration.mcpClient(), exportURI);
446+
447+
const localPathPart = contentWithExportPath(content);
448+
expect(localPathPart).toBeDefined();
449+
const [, localPath] = /"(.*)"/.exec(String(localPathPart?.text)) ?? [];
450+
expect(localPath).toBeDefined();
451+
452+
const exportedContent = JSON.parse(await fs.readFile(localPath as string, "utf8")) as Record<
453+
string,
454+
unknown
455+
>[];
456+
expect(exportedContent).toHaveLength(1);
457+
expect(exportedContent[0]?.name).toEqual("foo");
458+
});
374459
});
375460
},
376461
() => userConfig

0 commit comments

Comments
 (0)