Skip to content

Commit 8e62d2e

Browse files
blvablva
authored
Merge branch 'main' into MCP-68 
2 parents e570562 + 63597c2 commit 8e62d2e

File tree

6 files changed

+275
-71
lines changed

6 files changed

+275
-71
lines changed

‎.smithery/Dockerfile‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,10 @@
11
# Generated by https://smithery.ai. See: https://smithery.ai/docs/config#dockerfile
22
# ----- Build Stage -----
33
FROM node:lts-alpine AS builder
4+
5+
RUN adduser -D mcpuser
6+
USER mcpuser
7+
48
WORKDIR /app
59

610
# Copy package and configuration

‎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: 79 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -8,10 +8,16 @@ describeAccuracyTests([
88
{
99
toolName: "export",
1010
parameters: {
11+
exportTitle: Matcher.string(),
1112
database: "mflix",
1213
collection: "movies",
13-
filter: Matcher.emptyObjectOrUndefined,
14-
limit: Matcher.undefined,
14+
exportTitle: Matcher.string(),
15+
exportTarget: [
16+
{
17+
name: "find",
18+
arguments: {},
19+
},
20+
],
1521
},
1622
},
1723
],
@@ -22,11 +28,20 @@ describeAccuracyTests([
2228
{
2329
toolName: "export",
2430
parameters: {
31+
exportTitle: Matcher.string(),
2532
database: "mflix",
2633
collection: "movies",
27-
filter: {
28-
runtime: { $lt: 100 },
29-
},
34+
exportTitle: Matcher.string(),
35+
exportTarget: [
36+
{
37+
name: "find",
38+
arguments: {
39+
filter: {
40+
runtime: { $lt: 100 },
41+
},
42+
},
43+
},
44+
],
3045
},
3146
},
3247
],
@@ -37,16 +52,25 @@ describeAccuracyTests([
3752
{
3853
toolName: "export",
3954
parameters: {
55+
exportTitle: Matcher.string(),
4056
database: "mflix",
4157
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,
58+
exportTitle: Matcher.string(),
59+
exportTarget: [
60+
{
61+
name: "find",
62+
arguments: {
63+
projection: {
64+
title: 1,
65+
_id: Matcher.anyOf(
66+
Matcher.undefined,
67+
Matcher.number((value) => value === 0)
68+
),
69+
},
70+
filter: Matcher.emptyObjectOrUndefined,
71+
},
72+
},
73+
],
5074
},
5175
},
5276
],
@@ -57,11 +81,50 @@ describeAccuracyTests([
5781
{
5882
toolName: "export",
5983
parameters: {
84+
exportTitle: Matcher.string(),
6085
database: "mflix",
6186
collection: "movies",
62-
filter: { genres: "Horror" },
63-
sort: { runtime: 1 },
64-
limit: 2,
87+
exportTitle: Matcher.string(),
88+
exportTarget: [
89+
{
90+
name: "find",
91+
arguments: {
92+
filter: { genres: "Horror" },
93+
sort: { runtime: 1 },
94+
limit: 2,
95+
},
96+
},
97+
],
98+
},
99+
},
100+
],
101+
},
102+
{
103+
prompt: "Export an aggregation that groups all movie titles by the field release_year from mflix.movies",
104+
expectedToolCalls: [
105+
{
106+
toolName: "export",
107+
parameters: {
108+
database: "mflix",
109+
collection: "movies",
110+
exportTitle: Matcher.string(),
111+
exportTarget: [
112+
{
113+
name: "aggregate",
114+
arguments: {
115+
pipeline: [
116+
{
117+
$group: {
118+
_id: "$release_year",
119+
titles: {
120+
$push: "$title",
121+
},
122+
},
123+
},
124+
],
125+
},
126+
},
127+
],
65128
},
66129
},
67130
],

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