FranklinWH-Toolbox/service/cmd/mcp/tools/data/mongo.go

package data

import (
	"context"
	"encoding/json"
	"fmt"

	"github.com/mark3labs/mcp-go/mcp"
	"github.com/mark3labs/mcp-go/server"
	"go.mongodb.org/mongo-driver/bson"
	"go.mongodb.org/mongo-driver/mongo/options"

	"fwh-toolbox-server/cmd/mcp/tools/helper"
	"fwh-toolbox-server/database"
)

// mongoAllowedCollections 允许查询的 collection 白名单
var mongoAllowedCollections = []string{
	"site_info",
	"sim_card_manage",
	"batch_download",
	"site_analysis",
	"shutdown_statistics",
	"alarm_statistics",
	"sites_monitor_config",
	"user_config",
}

func registerMongoTools(s *server.MCPServer) {
	s.AddTool(
		mcp.NewTool("data_mongo_find",
			mcp.WithDescription("直接查询 MongoDB 集合数据（只读）。用于 service 层无法满足的灵活查询场景。"),
			mcp.WithString("collection", mcp.Required(), mcp.Description(
				fmt.Sprintf("集合名，可选值: %v", mongoAllowedCollections),
			)),
			mcp.WithString("filter", mcp.Description("查询条件，JSON 格式的 MongoDB filter（如 {\"sn\": \"xxx\"}），不传则查全部")),
			mcp.WithNumber("limit", mcp.Description("返回数量限制，默认10，最大100")),
			mcp.WithNumber("skip", mcp.Description("跳过数量（分页偏移），默认0")),
			mcp.WithString("sort", mcp.Description("排序条件，JSON 格式（如 {\"createdAt\": -1}）")),
		),
		handleMongoFind,
	)

	s.AddTool(
		mcp.NewTool("data_mongo_count",
			mcp.WithDescription("统计 MongoDB 集合中符合条件的文档数量"),
			mcp.WithString("collection", mcp.Required(), mcp.Description(
				fmt.Sprintf("集合名，可选值: %v", mongoAllowedCollections),
			)),
			mcp.WithString("filter", mcp.Description("查询条件，JSON 格式的 MongoDB filter，不传则统计全部")),
		),
		handleMongoCount,
	)
}

func handleMongoFind(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
	}

	// limit
	limit := int64(req.GetFloat("limit", 10))
	if limit > 100 {
		limit = 100
	}
	if limit < 1 {
		limit = 10
	}

	// skip
	skip := int64(req.GetFloat("skip", 0))
	if skip < 0 {
		skip = 0
	}

	// sort
	var sort bson.M
	sortStr := req.GetString("sort", "")
	if sortStr != "" {
		if err := json.Unmarshal([]byte(sortStr), &sort); err != nil {
			return helper.ParamInvalidError("sort", fmt.Sprintf("JSON 解析失败: %v", err)), nil
		}
	}

	// 查询
	coll := database.GetCollection(collection)
	findOpts := options.Find().SetLimit(limit).SetSkip(skip)
	if sort != nil {
		findOpts.SetSort(sort)
	}

	cursor, err := coll.Find(ctx, filter, findOpts)
	if err != nil {
		return helper.ServiceError("MongoDB查询", err), nil
	}
	defer cursor.Close(ctx)

	var results []bson.M
	if err := cursor.All(ctx, &results); err != nil {
		return helper.ServiceError("读取MongoDB结果", err), nil
	}

	if len(results) == 0 {
		return helper.EmptyResult("查询结果为空"), nil
	}

	return helper.JSONResult(results), nil
}

func handleMongoCount(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
	}

	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
	}

	coll := database.GetCollection(collection)
	count, err := coll.CountDocuments(ctx, filter)
	if err != nil {
		return helper.ServiceError("MongoDB统计", err), nil
	}

	return helper.JSONResult(map[string]any{
		"collection": collection,
		"count":      count,
	}), nil
}

func isAllowedCollection(name string) bool {
	for _, c := range mongoAllowedCollections {
		if c == name {
			return true
		}
	}
	return false
}