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{
	// 站点信息管理
	"SiteInfoManageData",
	"SiteInfoManageRecord",
	// SIM卡管理
	"SimCardManage",
	"SimCardMonthlyData",
	// 批量下载
	"BatchDownloadConfig",
	// 站点分析
	"SiteAnalysisConfig",
	// 统计
	"StatisticsShutdown",
	"StatisticsAlarm",
	// 站点监控
	"SiteMonitorConfig",
	// 自动升级
	"AutoUpdateConfigOption",
	"AutoUpdateConfig",
	// 注意：UserInfo（内部用户身份数据）与 AiChatSessions（历史对话自由文本，
	// 字段级脱敏无效）已从白名单移除，禁止通过 data_mongo_* 直查以防隐私泄露。
}

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.WithArray("fields", mcp.WithStringItems(),
				mcp.Description("字段投影：要返回的字段名数组（如 [\"sn\",\"status\"]），留空返回全部字段。用于大数据量场景下减少返回体积、节省 token。")),
			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,
	)

	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) {
	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)
	}

	// 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)
	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
	}

	// total：符合 filter 的文档总数（不受 limit/skip 影响），便于大数据量遍历时判断进度
	total, err := coll.CountDocuments(ctx, filter)
	if err != nil {
		return helper.ServiceError("统计MongoDB总数", err), 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) {
	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.JSONResultOrFileAuto(map[string]any{
		"collection": collection,
		"count":      count,
	}), 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 {
	for _, c := range mongoAllowedCollections {
		if c == name {
			return true
		}
	}
	return false
}