You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
 
 

140 lines
3.9 KiB

package vTool
import (
"context"
"github.com/mark3labs/mcp-go/mcp"
"github.com/mark3labs/mcp-go/server"
"mcp-go-db/vConfig"
"mcp-go-db/vService"
)
func Run(s *server.MCPServer, dbSvc *vService.DMDBService, config *vConfig.Config) {
// 添加数据库工具
dbTools := []mcp.Tool{
DBInfoTool("数据库配置信息"),
AllTableTool(config, "获取所有表"),
FetchDataTool(config, "获取表数据"),
AddCommentTool(config, "添加注释"),
AddColumnTool(config, "添加字段"),
DropColumnTool(config, "删除字段"),
}
// 注册工具处理程序
for _, tool := range dbTools {
s.AddTool(tool, func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
switch request.Params.Name {
case "数据库配置信息":
return DBInfo(config)
case "获取所有表":
return AllTable(ctx, request, dbSvc)
case "获取表数据":
return FetchData(ctx, request, dbSvc)
case "添加注释":
return AddComment(ctx, request, dbSvc)
case "添加字段":
return AddColumn(ctx, request, dbSvc)
case "删除字段":
return DropColumn(ctx, request, dbSvc)
}
return mcp.NewToolResultError("未知操作"), nil
})
}
}
func DBInfoTool(name string) mcp.Tool {
return mcp.NewTool(name, /* "数据库配置信息" */
mcp.WithDescription("数据库配置信息"),
)
}
func AddColumnTool(config *vConfig.Config, name string) mcp.Tool {
return mcp.NewTool(name, /* "添加字段" */
mcp.WithDescription("为表添加新字段"),
mcp.WithString("库名",
mcp.Required(),
mcp.Description("被操作的数据库"),
mcp.DefaultString(config.Schema),
),
mcp.WithString("表名",
mcp.Required(),
mcp.Description("被操作的数据库的表名"),
),
mcp.WithString("字段名",
mcp.Required(),
mcp.Description("被操作的数据库的表内的字段名"),
),
mcp.WithString("字段类型",
mcp.Required(),
mcp.Description("被操作的数据库的表内的字段类型"),
),
)
}
func DropColumnTool(config *vConfig.Config, name string) mcp.Tool {
return mcp.NewTool(name, /* "删除字段" */
mcp.WithDescription("从表删除字段"),
mcp.WithString("库名",
mcp.Required(),
mcp.Description("被操作的数据库"),
mcp.DefaultString(config.Schema),
),
mcp.WithString("表名",
mcp.Required(),
mcp.Description("被操作的数据库的表内的字段名"),
),
mcp.WithString("字段名",
mcp.Required(),
mcp.Description("被操作的数据库的表内的字段名"),
),
)
}
func AllTableTool(config *vConfig.Config, name string) mcp.Tool {
return mcp.NewTool(name, /* "获取所有表" */
mcp.WithDescription("从指定表获取数据"),
mcp.WithString("库名",
mcp.Required(),
mcp.Description("被操作的数据库"),
mcp.DefaultString(config.Schema),
),
)
}
func AddCommentTool(config *vConfig.Config, name string) mcp.Tool {
return mcp.NewTool(name, /* "添加注释" */
mcp.WithDescription("为表的字段添加新注释"),
mcp.WithString("库名",
mcp.Required(),
mcp.Description("被操作的数据库"),
mcp.DefaultString(config.Schema),
),
mcp.WithString("表名",
mcp.Required(),
mcp.Description("被操作的数据库的表名"),
),
mcp.WithString("字段名",
mcp.Required(),
mcp.Description("被操作的数据库的表内的字段名"),
),
)
}
func FetchDataTool(config *vConfig.Config, name string) mcp.Tool {
return mcp.NewTool(name, /* "获取表数据" */
mcp.WithDescription("从指定表获取数据"),
mcp.WithString("库名",
mcp.Required(),
mcp.Description("被操作的数据库"),
mcp.DefaultString(config.Schema),
),
mcp.WithString("表名",
mcp.Required(),
mcp.Description("被操作的数据库的表名"),
),
mcp.WithNumber("限制",
mcp.Required(),
mcp.Description("查询限制数量,避免全部查询,默认为 10"),
mcp.DefaultNumber(10),
),
)
}