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