|
|
@@ -4,6 +4,7 @@ import (
|
|
|
"context"
|
|
|
"encoding/json"
|
|
|
"fmt"
|
|
|
+ "time"
|
|
|
|
|
|
"github.com/mark3labs/mcp-go/mcp"
|
|
|
"github.com/mark3labs/mcp-go/server"
|
|
|
@@ -80,6 +81,57 @@ func registerMongoTools(s *server.MCPServer) {
|
|
|
),
|
|
|
handleMongoExport,
|
|
|
)
|
|
|
+
|
|
|
+ s.AddTool(
|
|
|
+ mcp.NewTool("data_mongo_aggregate",
|
|
|
+ mcp.WithDescription("对 MongoDB 集合执行聚合管道(aggregation pipeline,只读)并将结果流式导出到本地文件。"+
|
|
|
+ "支持 $match/$project/$group/$addFields/$lookup(跨集合关联/JOIN)/$sort/$limit/$unwind/$facet 等只读阶段,"+
|
|
|
+ "可在数据库侧一次性完成过滤+关联+聚合,避免多次取数与 Agent 侧拼接。"+
|
|
|
+ "结果边脱敏边写入 NDJSON 文件(每行一个 JSON 对象),仅返回 localPath+count+summary(不返回数据正文)。"+
|
|
|
+ "安全约束:collection 与 $lookup/$unionWith/$graphLookup 引用的集合都必须在白名单内;"+
|
|
|
+ "禁止 $out/$merge/$function/$accumulator/$where 等写入或执行代码的阶段。"),
|
|
|
+ mcp.WithString("collection", mcp.Required(), mcp.Description(
|
|
|
+ fmt.Sprintf("聚合起点集合名,可选值: %v", mongoAllowedCollections),
|
|
|
+ )),
|
|
|
+ mcp.WithString("pipeline", mcp.Required(), mcp.Description(
|
|
|
+ "聚合管道,JSON 数组字符串,每个元素是一个聚合阶段。"+
|
|
|
+ "例:[{\"$match\":{\"province\":\"Puerto Rico\"}},{\"$group\":{\"_id\":\"$province\",\"n\":{\"$sum\":1}}}]")),
|
|
|
+ mcp.WithNumber("max_docs", mcp.Description(fmt.Sprintf(
|
|
|
+ "最大输出条数上限,默认 %d,硬上限 %d。聚合结果超过上限会被截断并在返回中标注。", mongoAggregateDefaultMax, mongoExportHardMax))),
|
|
|
+ ),
|
|
|
+ handleMongoAggregate,
|
|
|
+ )
|
|
|
+}
|
|
|
+
|
|
|
+// mongoAggregateDefaultMax 聚合输出默认条数上限(软保护,可由 max_docs 覆盖)。
|
|
|
+const mongoAggregateDefaultMax = 50000
|
|
|
+
|
|
|
+// mongoAllowedStages 聚合管道允许的只读阶段白名单。
|
|
|
+// 明确排除写入/执行代码的阶段:$out、$merge(写集合)、$function、$accumulator、$where(执行 JS)。
|
|
|
+var mongoAllowedStages = map[string]bool{
|
|
|
+ "$match": true,
|
|
|
+ "$project": true,
|
|
|
+ "$addFields": true,
|
|
|
+ "$set": true,
|
|
|
+ "$unset": true,
|
|
|
+ "$group": true,
|
|
|
+ "$sort": true,
|
|
|
+ "$limit": true,
|
|
|
+ "$skip": true,
|
|
|
+ "$count": true,
|
|
|
+ "$unwind": true,
|
|
|
+ "$lookup": true,
|
|
|
+ "$graphLookup": true,
|
|
|
+ "$unionWith": true,
|
|
|
+ "$facet": true,
|
|
|
+ "$bucket": true,
|
|
|
+ "$bucketAuto": true,
|
|
|
+ "$sortByCount": true,
|
|
|
+ "$replaceRoot": true,
|
|
|
+ "$replaceWith": true,
|
|
|
+ "$sample": true,
|
|
|
+ "$redact": true,
|
|
|
+ "$setWindowFields": true,
|
|
|
}
|
|
|
|
|
|
// mongoExportDefaultMax 全量导出默认上限(防误导出全表的软保护,可由 max_docs 覆盖)
|
|
|
@@ -288,6 +340,184 @@ func handleMongoExport(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallT
|
|
|
return res, nil
|
|
|
}
|
|
|
|
|
|
+func handleMongoAggregate(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
|
|
|
+ }
|
|
|
+
|
|
|
+ // 解析 pipeline(JSON 数组)
|
|
|
+ pipelineStr, err := req.RequireString("pipeline")
|
|
|
+ if err != nil {
|
|
|
+ return helper.ParamError("pipeline"), nil
|
|
|
+ }
|
|
|
+ var pipeline []bson.M
|
|
|
+ if err := json.Unmarshal([]byte(pipelineStr), &pipeline); err != nil {
|
|
|
+ return helper.ParamInvalidError("pipeline", fmt.Sprintf("需要 JSON 数组,每个元素是一个聚合阶段: %v", err)), nil
|
|
|
+ }
|
|
|
+ if len(pipeline) == 0 {
|
|
|
+ return helper.ParamInvalidError("pipeline", "pipeline 不能为空"), nil
|
|
|
+ }
|
|
|
+
|
|
|
+ // 安全校验:阶段白名单 + 引用集合白名单(递归)
|
|
|
+ if verr := validatePipeline(pipeline); verr != "" {
|
|
|
+ return helper.ParamInvalidError("pipeline", verr), nil
|
|
|
+ }
|
|
|
+
|
|
|
+ // max_docs 输出上限
|
|
|
+ maxDocs := int64(req.GetFloat("max_docs", float64(mongoAggregateDefaultMax)))
|
|
|
+ if maxDocs < 1 {
|
|
|
+ maxDocs = mongoAggregateDefaultMax
|
|
|
+ }
|
|
|
+ if maxDocs > mongoExportHardMax {
|
|
|
+ maxDocs = mongoExportHardMax
|
|
|
+ }
|
|
|
+
|
|
|
+ // 聚合可能较重:设置超时 + allowDiskUse 防止内存溢出。
|
|
|
+ aggCtx, cancel := context.WithTimeout(ctx, 60*time.Second)
|
|
|
+ defer cancel()
|
|
|
+
|
|
|
+ coll := database.GetCollection(collection)
|
|
|
+ allowDisk := true
|
|
|
+ cursor, err := coll.Aggregate(aggCtx, pipeline, &options.AggregateOptions{AllowDiskUse: &allowDisk})
|
|
|
+ if err != nil {
|
|
|
+ return helper.ServiceError("MongoDB聚合", err), nil
|
|
|
+ }
|
|
|
+ defer cursor.Close(aggCtx)
|
|
|
+
|
|
|
+ // 流式:逐条 decode → 脱敏 → 写 NDJSON,内存常数级;达到 maxDocs 即停止。
|
|
|
+ var emitted int64
|
|
|
+ res := helper.ExportStreamNDJSON(func() (any, bool, error) {
|
|
|
+ if emitted >= maxDocs {
|
|
|
+ return nil, false, nil // 达到上限,正常结束(截断)
|
|
|
+ }
|
|
|
+ if !cursor.Next(aggCtx) {
|
|
|
+ return nil, false, cursor.Err()
|
|
|
+ }
|
|
|
+ var doc bson.M
|
|
|
+ if err := cursor.Decode(&doc); err != nil {
|
|
|
+ return nil, false, err
|
|
|
+ }
|
|
|
+ emitted++
|
|
|
+ return doc, true, nil
|
|
|
+ }, "mongo_aggregate_"+collection)
|
|
|
+ return res, nil
|
|
|
+}
|
|
|
+
|
|
|
+// validatePipeline 对聚合管道做安全校验,返回空字符串表示通过,否则返回错误说明。
|
|
|
+// - 每个阶段的顶层操作符必须在 mongoAllowedStages 白名单内(排除 $out/$merge/$function 等)。
|
|
|
+// - 递归收集 $lookup/$graphLookup/$unionWith 引用的集合,全部必须在集合白名单内,
|
|
|
+// 防止借跨集合阶段绕过白名单读取被禁集合。
|
|
|
+func validatePipeline(pipeline []bson.M) string {
|
|
|
+ for i, stage := range pipeline {
|
|
|
+ if len(stage) != 1 {
|
|
|
+ return fmt.Sprintf("第 %d 个阶段必须恰好包含一个聚合操作符", i+1)
|
|
|
+ }
|
|
|
+ for op, spec := range stage {
|
|
|
+ if !mongoAllowedStages[op] {
|
|
|
+ return fmt.Sprintf("第 %d 个阶段使用了不允许的操作符 %q(仅允许只读阶段,禁止 $out/$merge/$function 等)", i+1, op)
|
|
|
+ }
|
|
|
+ if refs := collectStageCollectionRefs(op, spec); len(refs) > 0 {
|
|
|
+ for _, c := range refs {
|
|
|
+ if !isAllowedCollection(c) {
|
|
|
+ return fmt.Sprintf("第 %d 个阶段 %s 引用了不在白名单内的集合 %q,可选: %v", i+1, op, c, mongoAllowedCollections)
|
|
|
+ }
|
|
|
+ }
|
|
|
+ }
|
|
|
+ // $unionWith 可能携带嵌套 pipeline,需递归校验
|
|
|
+ if op == "$unionWith" {
|
|
|
+ if sub := extractUnionWithPipeline(spec); sub != nil {
|
|
|
+ if verr := validatePipeline(sub); verr != "" {
|
|
|
+ return fmt.Sprintf("第 %d 个阶段 $unionWith 的嵌套 pipeline 校验失败: %s", i+1, verr)
|
|
|
+ }
|
|
|
+ }
|
|
|
+ }
|
|
|
+ // $lookup 的子 pipeline 形式(带 pipeline 字段)也需递归校验
|
|
|
+ if op == "$lookup" {
|
|
|
+ if sub := extractLookupPipeline(spec); sub != nil {
|
|
|
+ if verr := validatePipeline(sub); verr != "" {
|
|
|
+ return fmt.Sprintf("第 %d 个阶段 $lookup 的嵌套 pipeline 校验失败: %s", i+1, verr)
|
|
|
+ }
|
|
|
+ }
|
|
|
+ }
|
|
|
+ }
|
|
|
+ }
|
|
|
+ return ""
|
|
|
+}
|
|
|
+
|
|
|
+// collectStageCollectionRefs 提取某个阶段引用的外部集合名(用于白名单校验)。
|
|
|
+// - $lookup: spec.from(字符串)
|
|
|
+// - $graphLookup: spec.from(字符串)
|
|
|
+// - $unionWith: 字符串本身,或 spec.coll(字符串)
|
|
|
+func collectStageCollectionRefs(op string, spec any) []string {
|
|
|
+ var refs []string
|
|
|
+ switch op {
|
|
|
+ case "$lookup", "$graphLookup":
|
|
|
+ if m, ok := toMap(spec); ok {
|
|
|
+ if from, ok := m["from"].(string); ok && from != "" {
|
|
|
+ refs = append(refs, from)
|
|
|
+ }
|
|
|
+ }
|
|
|
+ case "$unionWith":
|
|
|
+ if s, ok := spec.(string); ok && s != "" {
|
|
|
+ refs = append(refs, s)
|
|
|
+ } else if m, ok := toMap(spec); ok {
|
|
|
+ if coll, ok := m["coll"].(string); ok && coll != "" {
|
|
|
+ refs = append(refs, coll)
|
|
|
+ }
|
|
|
+ }
|
|
|
+ }
|
|
|
+ return refs
|
|
|
+}
|
|
|
+
|
|
|
+// extractUnionWithPipeline 提取 $unionWith 的嵌套 pipeline(若有)。
|
|
|
+func extractUnionWithPipeline(spec any) []bson.M {
|
|
|
+ m, ok := toMap(spec)
|
|
|
+ if !ok {
|
|
|
+ return nil
|
|
|
+ }
|
|
|
+ return toPipeline(m["pipeline"])
|
|
|
+}
|
|
|
+
|
|
|
+// extractLookupPipeline 提取 $lookup 的子 pipeline(若有)。
|
|
|
+func extractLookupPipeline(spec any) []bson.M {
|
|
|
+ m, ok := toMap(spec)
|
|
|
+ if !ok {
|
|
|
+ return nil
|
|
|
+ }
|
|
|
+ return toPipeline(m["pipeline"])
|
|
|
+}
|
|
|
+
|
|
|
+// toMap 将 bson.M / map[string]any 统一转为 map[string]any。
|
|
|
+func toMap(v any) (map[string]any, bool) {
|
|
|
+ switch t := v.(type) {
|
|
|
+ case bson.M:
|
|
|
+ return map[string]any(t), true
|
|
|
+ case map[string]any:
|
|
|
+ return t, true
|
|
|
+ default:
|
|
|
+ return nil, false
|
|
|
+ }
|
|
|
+}
|
|
|
+
|
|
|
+// toPipeline 将任意值尽力转换为 []bson.M(聚合子管道)。
|
|
|
+func toPipeline(v any) []bson.M {
|
|
|
+ arr, ok := v.([]any)
|
|
|
+ if !ok {
|
|
|
+ return nil
|
|
|
+ }
|
|
|
+ out := make([]bson.M, 0, len(arr))
|
|
|
+ for _, e := range arr {
|
|
|
+ if m, ok := toMap(e); ok {
|
|
|
+ out = append(out, bson.M(m))
|
|
|
+ }
|
|
|
+ }
|
|
|
+ return out
|
|
|
+}
|
|
|
+
|
|
|
func isAllowedCollection(name string) bool {
|
|
|
for _, c := range mongoAllowedCollections {
|
|
|
if c == name {
|