|
@@ -46,6 +46,8 @@ func registerMongoTools(s *server.MCPServer) {
|
|
|
fmt.Sprintf("集合名,可选值: %v", mongoAllowedCollections),
|
|
fmt.Sprintf("集合名,可选值: %v", mongoAllowedCollections),
|
|
|
)),
|
|
)),
|
|
|
mcp.WithString("filter", mcp.Description("查询条件,JSON 格式的 MongoDB filter(如 {\"sn\": \"xxx\"}),不传则查全部")),
|
|
mcp.WithString("filter", mcp.Description("查询条件,JSON 格式的 MongoDB filter(如 {\"sn\": \"xxx\"}),不传则查全部")),
|
|
|
|
|
+ mcp.WithArray("fields", mcp.WithStringItems(),
|
|
|
|
|
+ mcp.Description("字段投影:要返回的字段名数组(如 [\"sn\",\"status\"]),留空返回全部字段。用于大数据量场景下减少返回体积、节省 token。")),
|
|
|
mcp.WithNumber("limit", mcp.Description("返回数量限制,默认10,最大100")),
|
|
mcp.WithNumber("limit", mcp.Description("返回数量限制,默认10,最大100")),
|
|
|
mcp.WithNumber("skip", mcp.Description("跳过数量(分页偏移),默认0")),
|
|
mcp.WithNumber("skip", mcp.Description("跳过数量(分页偏移),默认0")),
|
|
|
mcp.WithString("sort", mcp.Description("排序条件,JSON 格式(如 {\"createdAt\": -1})")),
|
|
mcp.WithString("sort", mcp.Description("排序条件,JSON 格式(如 {\"createdAt\": -1})")),
|
|
@@ -63,8 +65,30 @@ func registerMongoTools(s *server.MCPServer) {
|
|
|
),
|
|
),
|
|
|
handleMongoCount,
|
|
handleMongoCount,
|
|
|
)
|
|
)
|
|
|
|
|
+
|
|
|
|
|
+ s.AddTool(
|
|
|
|
|
+ mcp.NewTool("data_mongo_export",
|
|
|
|
|
+ mcp.WithDescription("全量导出 MongoDB 集合数据到本地文件(只读)。用于大数据量导出场景:不受 find 的 100 条上限限制,游标流式拉取符合条件的全部文档,边脱敏边写入 NDJSON 文件(每行一个 JSON 对象),内存占用为常数级。仅返回 localPath、count 与结构化摘要(不返回数据正文,避免占用上下文)。为防止误导出全表,无 filter 时若总数超过上限会被拒绝。"),
|
|
|
|
|
+ mcp.WithString("collection", mcp.Required(), mcp.Description(
|
|
|
|
|
+ fmt.Sprintf("集合名,可选值: %v", mongoAllowedCollections),
|
|
|
|
|
+ )),
|
|
|
|
|
+ mcp.WithString("filter", mcp.Description("查询条件,JSON 格式的 MongoDB filter,不传则导出全部(受上限保护)")),
|
|
|
|
|
+ mcp.WithArray("fields", mcp.WithStringItems(),
|
|
|
|
|
+ mcp.Description("字段投影:只导出这些字段(如 [\"sn\",\"status\"]),留空导出全部字段。")),
|
|
|
|
|
+ mcp.WithString("sort", mcp.Description("排序条件,JSON 格式(如 {\"createdAt\": -1})")),
|
|
|
|
|
+ mcp.WithNumber("max_docs", mcp.Description(fmt.Sprintf("最大导出条数上限,默认 %d,硬上限 %d。超过则拒绝,请缩小 filter。", mongoExportDefaultMax, mongoExportHardMax))),
|
|
|
|
|
+ ),
|
|
|
|
|
+ handleMongoExport,
|
|
|
|
|
+ )
|
|
|
}
|
|
}
|
|
|
|
|
|
|
|
|
|
+// mongoExportDefaultMax 全量导出默认上限(防误导出全表的软保护,可由 max_docs 覆盖)
|
|
|
|
|
+const mongoExportDefaultMax = 50000
|
|
|
|
|
+
|
|
|
|
|
+// mongoExportHardMax 全量导出硬上限。流式导出内存为常数级,此上限仅用于防止
|
|
|
|
|
+// 误导出超大规模数据,而非内存保护,故设置较高。
|
|
|
|
|
+const mongoExportHardMax = 5000000
|
|
|
|
|
+
|
|
|
func handleMongoFind(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
|
|
func handleMongoFind(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
|
|
|
collection, err := req.RequireString("collection")
|
|
collection, err := req.RequireString("collection")
|
|
|
if err != nil {
|
|
if err != nil {
|
|
@@ -113,6 +137,15 @@ func handleMongoFind(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToo
|
|
|
findOpts.SetSort(sort)
|
|
findOpts.SetSort(sort)
|
|
|
}
|
|
}
|
|
|
|
|
|
|
|
|
|
+ // fields 字段投影:留空返回全部字段;指定则只返回这些字段(源头减量,节省 token)
|
|
|
|
|
+ if fields, ok := helper.GetStringSlice(req, "fields"); ok && len(fields) > 0 {
|
|
|
|
|
+ projection := bson.M{}
|
|
|
|
|
+ for _, f := range fields {
|
|
|
|
|
+ projection[f] = 1
|
|
|
|
|
+ }
|
|
|
|
|
+ findOpts.SetProjection(projection)
|
|
|
|
|
+ }
|
|
|
|
|
+
|
|
|
cursor, err := coll.Find(ctx, filter, findOpts)
|
|
cursor, err := coll.Find(ctx, filter, findOpts)
|
|
|
if err != nil {
|
|
if err != nil {
|
|
|
return helper.ServiceError("MongoDB查询", err), nil
|
|
return helper.ServiceError("MongoDB查询", err), nil
|
|
@@ -124,11 +157,24 @@ func handleMongoFind(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToo
|
|
|
return helper.ServiceError("读取MongoDB结果", err), nil
|
|
return helper.ServiceError("读取MongoDB结果", err), nil
|
|
|
}
|
|
}
|
|
|
|
|
|
|
|
- if len(results) == 0 {
|
|
|
|
|
- return helper.EmptyResult("查询结果为空"), nil
|
|
|
|
|
|
|
+ // total:符合 filter 的文档总数(不受 limit/skip 影响),便于大数据量遍历时判断进度
|
|
|
|
|
+ total, err := coll.CountDocuments(ctx, filter)
|
|
|
|
|
+ if err != nil {
|
|
|
|
|
+ return helper.ServiceError("统计MongoDB总数", err), nil
|
|
|
}
|
|
}
|
|
|
|
|
|
|
|
- return helper.JSONResultOrFileAuto(results), nil
|
|
|
|
|
|
|
+ returned := int64(len(results))
|
|
|
|
|
+ hasMore := skip+returned < total
|
|
|
|
|
+
|
|
|
|
|
+ // 结构化返回:附带 total/returned/skip/limit/hasMore,便于 AI 判断是否翻完、如何翻页
|
|
|
|
|
+ return helper.JSONResultOrFileAuto(map[string]any{
|
|
|
|
|
+ "total": total,
|
|
|
|
|
+ "returned": returned,
|
|
|
|
|
+ "skip": skip,
|
|
|
|
|
+ "limit": limit,
|
|
|
|
|
+ "hasMore": hasMore,
|
|
|
|
|
+ "items": results,
|
|
|
|
|
+ }), nil
|
|
|
}
|
|
}
|
|
|
|
|
|
|
|
func handleMongoCount(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
|
|
func handleMongoCount(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
|
|
@@ -159,6 +205,89 @@ func handleMongoCount(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallTo
|
|
|
}), nil
|
|
}), nil
|
|
|
}
|
|
}
|
|
|
|
|
|
|
|
|
|
+func handleMongoExport(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
|
|
|
|
|
+ collection, err := req.RequireString("collection")
|
|
|
|
|
+ if err != nil {
|
|
|
|
|
+ return helper.ParamError("collection"), nil
|
|
|
|
|
+ }
|
|
|
|
|
+ if !isAllowedCollection(collection) {
|
|
|
|
|
+ return helper.ParamInvalidError("collection", fmt.Sprintf("不允许查询此集合,可选: %v", mongoAllowedCollections)), nil
|
|
|
|
|
+ }
|
|
|
|
|
+
|
|
|
|
|
+ // 解析 filter
|
|
|
|
|
+ filterStr := req.GetString("filter", "{}")
|
|
|
|
|
+ var filter bson.M
|
|
|
|
|
+ if err := json.Unmarshal([]byte(filterStr), &filter); err != nil {
|
|
|
|
|
+ return helper.ParamInvalidError("filter", fmt.Sprintf("JSON 解析失败: %v", err)), nil
|
|
|
|
|
+ }
|
|
|
|
|
+ hasFilter := len(filter) > 0
|
|
|
|
|
+
|
|
|
|
|
+ // max_docs 上限
|
|
|
|
|
+ maxDocs := int64(req.GetFloat("max_docs", float64(mongoExportDefaultMax)))
|
|
|
|
|
+ if maxDocs < 1 {
|
|
|
|
|
+ maxDocs = mongoExportDefaultMax
|
|
|
|
|
+ }
|
|
|
|
|
+ if maxDocs > mongoExportHardMax {
|
|
|
|
|
+ maxDocs = mongoExportHardMax
|
|
|
|
|
+ }
|
|
|
|
|
+
|
|
|
|
|
+ coll := database.GetCollection(collection)
|
|
|
|
|
+
|
|
|
|
|
+ // 先统计总量,做防误导出保护
|
|
|
|
|
+ total, err := coll.CountDocuments(ctx, filter)
|
|
|
|
|
+ if err != nil {
|
|
|
|
|
+ return helper.ServiceError("统计导出总数", err), nil
|
|
|
|
|
+ }
|
|
|
|
|
+ if total == 0 {
|
|
|
|
|
+ return helper.EmptyResult("符合条件的文档为 0,无需导出"), nil
|
|
|
|
|
+ }
|
|
|
|
|
+ if total > maxDocs {
|
|
|
|
|
+ hint := "请缩小 filter 或调高 max_docs"
|
|
|
|
|
+ if !hasFilter {
|
|
|
|
|
+ hint = "未提供 filter 且数据量过大,请提供 filter 缩小范围,或显式调高 max_docs"
|
|
|
|
|
+ }
|
|
|
|
|
+ return helper.ParamInvalidError("max_docs", fmt.Sprintf(
|
|
|
|
|
+ "符合条件的文档数 %d 超过上限 %d,%s(硬上限 %d)", total, maxDocs, hint, mongoExportHardMax)), nil
|
|
|
|
|
+ }
|
|
|
|
|
+
|
|
|
|
|
+ // 组装查询选项
|
|
|
|
|
+ findOpts := options.Find()
|
|
|
|
|
+ if sortStr := req.GetString("sort", ""); sortStr != "" {
|
|
|
|
|
+ var sort bson.M
|
|
|
|
|
+ if err := json.Unmarshal([]byte(sortStr), &sort); err != nil {
|
|
|
|
|
+ return helper.ParamInvalidError("sort", fmt.Sprintf("JSON 解析失败: %v", err)), nil
|
|
|
|
|
+ }
|
|
|
|
|
+ findOpts.SetSort(sort)
|
|
|
|
|
+ }
|
|
|
|
|
+ if fields, ok := helper.GetStringSlice(req, "fields"); ok && len(fields) > 0 {
|
|
|
|
|
+ projection := bson.M{}
|
|
|
|
|
+ for _, f := range fields {
|
|
|
|
|
+ projection[f] = 1
|
|
|
|
|
+ }
|
|
|
|
|
+ findOpts.SetProjection(projection)
|
|
|
|
|
+ }
|
|
|
|
|
+
|
|
|
|
|
+ // 游标流式导出:逐条 decode → 脱敏 → 写 NDJSON,内存常数级
|
|
|
|
|
+ cursor, err := coll.Find(ctx, filter, findOpts)
|
|
|
|
|
+ if err != nil {
|
|
|
|
|
+ return helper.ServiceError("MongoDB导出查询", err), nil
|
|
|
|
|
+ }
|
|
|
|
|
+ defer cursor.Close(ctx)
|
|
|
|
|
+
|
|
|
|
|
+ res := helper.ExportStreamNDJSON(func() (any, bool, error) {
|
|
|
|
|
+ if !cursor.Next(ctx) {
|
|
|
|
|
+ // Next 返回 false:可能是正常结束,也可能是出错
|
|
|
|
|
+ return nil, false, cursor.Err()
|
|
|
|
|
+ }
|
|
|
|
|
+ var doc bson.M
|
|
|
|
|
+ if err := cursor.Decode(&doc); err != nil {
|
|
|
|
|
+ return nil, false, err
|
|
|
|
|
+ }
|
|
|
|
|
+ return doc, true, nil
|
|
|
|
|
+ }, "mongo_export_"+collection)
|
|
|
|
|
+ return res, nil
|
|
|
|
|
+}
|
|
|
|
|
+
|
|
|
func isAllowedCollection(name string) bool {
|
|
func isAllowedCollection(name string) bool {
|
|
|
for _, c := range mongoAllowedCollections {
|
|
for _, c := range mongoAllowedCollections {
|
|
|
if c == name {
|
|
if c == name {
|