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

package data

import (
	"context"
	"fmt"
	"time"

	"github.com/mark3labs/mcp-go/mcp"
	"github.com/mark3labs/mcp-go/server"

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

func registerPgsqlTools(s *server.MCPServer) {
	s.AddTool(
		mcp.NewTool("data_pgsql_query_snapshot_file_record",
			mcp.WithDescription("查询 PostgreSQL 中的告警快照文件解析记录"),
			mcp.WithString("file_name", mcp.Required(), mcp.Description("快照文件名")),
		),
		handlePgsqlQuerySnapshotFileRecord,
	)

	s.AddTool(
		mcp.NewTool("data_pgsql_query_snapshot_data",
			mcp.WithDescription("查询 PostgreSQL 中的告警快照数据（按SN和日期）"),
			mcp.WithString("sn", mcp.Required(), mcp.Description("站点SN")),
			mcp.WithString("date", mcp.Required(), mcp.Description("日期，格式 yyyy-MM-dd")),
		),
		handlePgsqlQuerySnapshotData,
	)
}

func handlePgsqlQuerySnapshotFileRecord(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
	fileName, err := req.RequireString("file_name")
	if err != nil {
		return helper.ParamError("file_name"), nil
	}

	data, err := repositories.QuerySnapshotFileRecord(ctx, fileName)
	if err != nil {
		return helper.EmptyResult(fmt.Sprintf("未找到文件记录: %v", err)), nil
	}
	return helper.JSONResult(data), nil
}

func handlePgsqlQuerySnapshotData(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
	sn, err := req.RequireString("sn")
	if err != nil {
		return helper.ParamError("sn"), nil
	}
	dateStr, err := req.RequireString("date")
	if err != nil {
		return helper.ParamError("date"), nil
	}

	// 验证日期格式
	if _, err := time.Parse("2006-01-02", dateStr); err != nil {
		return helper.ParamInvalidError("date", fmt.Sprintf("需要 yyyy-MM-dd 格式: %v", err)), nil
	}

	data, err := repositories.QuerySnapshotData(ctx, sn, dateStr)
	if err != nil {
		return helper.ServiceError("查询快照数据", err), nil
	}

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