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.
 
 

522 lines
18 KiB

package vTool
import (
"context"
"encoding/csv"
"encoding/json"
"fmt"
"github.com/mark3labs/mcp-go/mcp"
"github.com/mark3labs/mcp-go/server"
"mcp-go-db/vConfig"
"mcp-go-db/vService"
"strings"
)
// RegisterDBTools 注册数据库工具
func RegisterDBTools(s *server.MCPServer, dbSvc *vService.DMDBService, config *vConfig.Config) {
// 添加数据库工具
dbTools := []mcp.Tool{
DBInfoTool("数据库配置信息"),
AllTableTool(config, "获取所有表"),
FetchDataTool(config, "获取表数据"),
TableStructureTool(config, "表结构详情"),
AddCommentTool(config, "添加注释"),
AddColumnTool(config, "添加字段"),
DropColumnTool(config, "删除字段"),
InsertRowTool(config, "插入记录"),
UpdateRowsTool(config, "更新记录"),
DeleteRowsTool(config, "删除记录"),
CreateTableTool(config, "创建表"),
DropTableTool(config, "删除表"),
SQLTool(config, "执行SQL"),
}
// 注册工具处理程序
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 TableStructure(ctx, request, dbSvc)
case "添加注释":
return AddComment(ctx, request, dbSvc)
case "添加字段":
return AddColumn(ctx, request, dbSvc)
case "删除字段":
return DropColumn(ctx, request, dbSvc)
case "插入记录":
return InsertRow(ctx, request, dbSvc)
case "更新记录":
return UpdateRows(ctx, request, dbSvc)
case "删除记录":
return DeleteRows(ctx, request, dbSvc)
case "创建表":
return CreateTable(ctx, request, dbSvc)
case "删除表":
return DropTable(ctx, request, dbSvc)
case "执行SQL":
return ExecuteSQL(ctx, request, dbSvc)
}
return mcp.NewToolResultError("未知操作"), nil
})
}
}
func DBInfoTool(name string) mcp.Tool {
return mcp.NewTool(name, /* "数据库配置信息" */
mcp.WithDescription("数据库配置信息"),
)
}
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 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("被操作的数据库的表内的字段名"),
),
)
}
// InsertRowTool 插入记录
func InsertRowTool(config *vConfig.Config, name string) mcp.Tool {
return mcp.NewTool(name,
mcp.WithDescription("向表中插入一行记录,键值对为 JSON 对象,如 {\"NAME\":\"张三\",\"AGE\":20}"),
mcp.WithString("库名", mcp.Required(), mcp.Description("库名"), mcp.DefaultString(config.Schema)),
mcp.WithString("表名", mcp.Required(), mcp.Description("表名")),
mcp.WithString("键值对", mcp.Required(), mcp.Description("JSON 对象,键为列名,值为插入值")),
)
}
// UpdateRowsTool 更新记录
func UpdateRowsTool(config *vConfig.Config, name string) mcp.Tool {
return mcp.NewTool(name,
mcp.WithDescription("按条件更新表记录,SET 与 WHERE 使用占位符 ?,值按顺序放在「值列表」JSON 数组中"),
mcp.WithString("库名", mcp.Required(), mcp.Description("库名"), mcp.DefaultString(config.Schema)),
mcp.WithString("表名", mcp.Required(), mcp.Description("表名")),
mcp.WithString("SET子句", mcp.Required(), mcp.Description("例如 NAME=?, AGE=?")),
mcp.WithString("WHERE子句", mcp.Required(), mcp.Description("例如 ID=?")),
mcp.WithString("值列表", mcp.Required(), mcp.Description("JSON 数组,与 SET 和 WHERE 中 ? 顺序一致,如 [\"张三\",20,1]")),
)
}
// DeleteRowsTool 删除记录
func DeleteRowsTool(config *vConfig.Config, name string) mcp.Tool {
return mcp.NewTool(name,
mcp.WithDescription("按条件删除表记录,WHERE 使用占位符 ?,值放在「值列表」JSON 数组中"),
mcp.WithString("库名", mcp.Required(), mcp.Description("库名"), mcp.DefaultString(config.Schema)),
mcp.WithString("表名", mcp.Required(), mcp.Description("表名")),
mcp.WithString("WHERE子句", mcp.Required(), mcp.Description("例如 ID=?")),
mcp.WithString("值列表", mcp.Required(), mcp.Description("JSON 数组,与 WHERE 中 ? 顺序一致,如 [1]")),
)
}
// CreateTableTool 创建表
func CreateTableTool(config *vConfig.Config, name string) mcp.Tool {
return mcp.NewTool(name,
mcp.WithDescription("创建表,字段定义如:ID INT PRIMARY KEY, NAME VARCHAR(100)"),
mcp.WithString("库名", mcp.Required(), mcp.Description("库名"), mcp.DefaultString(config.Schema)),
mcp.WithString("表名", mcp.Required(), mcp.Description("表名")),
mcp.WithString("字段定义", mcp.Required(), mcp.Description("列定义,逗号分隔,如 ID INT, NAME VARCHAR(100)")),
)
}
// DropTableTool 删除表
func DropTableTool(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("表名")),
)
}
// SQLTool 创建SQL执行工具
func SQLTool(config *vConfig.Config, name string) mcp.Tool {
return mcp.NewTool(name,
mcp.WithDescription("执行SQL语句"),
mcp.WithString("sql",
mcp.Required(),
mcp.Description("要执行的SQL语句"),
),
mcp.WithString("库名",
mcp.Required(),
mcp.Description("目标数据库"),
mcp.DefaultString(config.Schema),
),
)
}
// ExecuteSQL 执行SQL语句
func ExecuteSQL(ctx context.Context, request mcp.CallToolRequest, dbSvc *vService.DMDBService) (*mcp.CallToolResult, error) {
sql := request.Params.Arguments["sql"].(string)
dbName := request.Params.Arguments["库名"].(string)
// 检查SQL类型
upperSQL := strings.ToUpper(strings.TrimSpace(sql))
isDDL := strings.HasPrefix(upperSQL, "CREATE") || strings.HasPrefix(upperSQL, "ALTER") ||
strings.HasPrefix(upperSQL, "DROP") || strings.HasPrefix(upperSQL, "TRUNCATE")
isDML := strings.HasPrefix(upperSQL, "INSERT") || strings.HasPrefix(upperSQL, "UPDATE") ||
strings.HasPrefix(upperSQL, "DELETE")
// 需要确认的操作
if isDDL || isDML {
return &mcp.CallToolResult{
Content: []mcp.Content{
mcp.NewTextContent("请确认执行以下SQL:\n" + sql + "\n\n类型: " +
func() string {
if isDDL {
return "DDL(结构变更)"
}
return "DML(数据变更)"
}()),
},
}, nil
}
// 执行查询
results, err := dbSvc.QueryBySQL(ctx, dbName, sql, 50) // 限制50条记录
if err != nil {
return mcp.NewToolResultError(err.Error()), nil
}
// 转换为CSV格式
var csvData strings.Builder
writer := csv.NewWriter(&csvData)
for _, row := range results {
if err := writer.Write(row); err != nil {
return mcp.NewToolResultError("CSV转换失败"), nil
}
}
writer.Flush()
return &mcp.CallToolResult{
Content: []mcp.Content{
mcp.NewTextContent(csvData.String()),
},
}, nil
}
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("被操作的数据库的表内的字段名"),
),
mcp.WithString("注释描述",
mcp.Required(),
mcp.Description("要添加的注释内容"),
),
)
}
func TableStructureTool(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("表名"),
),
)
}
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),
),
)
}
// DBInfo 查询系统环境信息
func DBInfo(config *vConfig.Config) (*mcp.CallToolResult, error) {
return mcp.NewToolResultText("环境为:" + config.Database.DmDsn), nil
}
// AllTable 查询数据Tool
func AllTable(ctx context.Context, request mcp.CallToolRequest, dbSvc *vService.DMDBService) (*mcp.CallToolResult, error) {
dbName, _ := request.Params.Arguments["库名"].(string)
data, err := dbSvc.AllTable(ctx, dbName)
if err != nil {
return mcp.NewToolResultError(err.Error()), nil
}
jsonData, err := json.Marshal(data)
if err != nil {
// 错误处理
return mcp.NewToolResultError("JSON 序列化失败"), nil
}
return mcp.NewToolResultText(string(jsonData)), nil
}
// TableStructure 表结构详情
func TableStructure(ctx context.Context, request mcp.CallToolRequest, dbSvc *vService.DMDBService) (*mcp.CallToolResult, error) {
dbName, _ := request.Params.Arguments["库名"].(string)
tableName, _ := request.Params.Arguments["表名"].(string)
list, err := dbSvc.TableStructure(ctx, dbName, tableName)
if err != nil {
return mcp.NewToolResultError(err.Error()), nil
}
jsonData, err := json.Marshal(list)
if err != nil {
return mcp.NewToolResultError("JSON序列化失败"), nil
}
return mcp.NewToolResultText(string(jsonData)), nil
}
// FetchData 查询数据Tool
func FetchData(ctx context.Context, request mcp.CallToolRequest, dbSvc *vService.DMDBService) (*mcp.CallToolResult, error) {
dbName, ok := request.Params.Arguments["库名"].(string)
tableName, ok := request.Params.Arguments["表名"].(string)
if !ok {
return mcp.NewToolResultError("缺少表名参数"), nil
}
limit := 10
if limitVal, ok := request.Params.Arguments["限制"].(float64); ok {
limit = int(limitVal)
}
data, err := dbSvc.FetchData(ctx, dbName, tableName, limit)
if err != nil {
return mcp.NewToolResultError(err.Error()), nil
}
jsonData, err := json.Marshal(data)
if err != nil {
// 错误处理
return mcp.NewToolResultError("JSON 序列化失败"), nil
}
return mcp.NewToolResultText(string(jsonData)), nil
}
// AddColumn 添加字段Tool
func AddColumn(ctx context.Context, request mcp.CallToolRequest, dbSvc *vService.DMDBService) (*mcp.CallToolResult, error) {
dbName, ok := request.Params.Arguments["库名"].(string)
tableName, ok := request.Params.Arguments["表名"].(string)
if !ok {
return mcp.NewToolResultError("缺少表名参数"), nil
}
columnName, ok := request.Params.Arguments["字段名"].(string)
if !ok {
return mcp.NewToolResultError("缺少字段名参数"), nil
}
columnType, ok := request.Params.Arguments["字段类型"].(string)
if !ok {
return mcp.NewToolResultError("缺少字段类型参数"), nil
}
if err := dbSvc.AddColumn(ctx, dbName, tableName, columnName, columnType); err != nil {
return mcp.NewToolResultError(err.Error()), nil
}
return mcp.NewToolResultText("字段添加成功"), nil
}
// AddComment 添加注释Tool
func AddComment(ctx context.Context, request mcp.CallToolRequest, dbSvc *vService.DMDBService) (*mcp.CallToolResult, error) {
dbName, ok := request.Params.Arguments["库名"].(string)
tableName, ok := request.Params.Arguments["表名"].(string)
if !ok {
return mcp.NewToolResultError("缺少表名参数"), nil
}
columnName, ok := request.Params.Arguments["字段名"].(string)
if !ok {
return mcp.NewToolResultError("缺少字段名参数"), nil
}
columnComment, ok := request.Params.Arguments["注释描述"].(string)
if !ok {
return mcp.NewToolResultError("缺少注释参数"), nil
}
if err := dbSvc.AddComment(ctx, dbName, tableName, columnName, columnComment); err != nil {
return mcp.NewToolResultError(err.Error()), nil
}
return mcp.NewToolResultText("字段添加注释成功"), nil
}
// DropColumn 字段删除Tool
func DropColumn(ctx context.Context, request mcp.CallToolRequest, dbSvc *vService.DMDBService) (*mcp.CallToolResult, error) {
dbName, ok := request.Params.Arguments["库名"].(string)
tableName, ok := request.Params.Arguments["表名"].(string)
if !ok {
return mcp.NewToolResultError("缺少表名参数"), nil
}
columnName, ok := request.Params.Arguments["字段名"].(string)
if !ok {
return mcp.NewToolResultError("缺少字段名"), nil
}
if err := dbSvc.DropColumn(ctx, dbName, tableName, columnName); err != nil {
return mcp.NewToolResultError(err.Error()), nil
}
return mcp.NewToolResultText("字段删除成功"), nil
}
// InsertRow 插入记录
func InsertRow(ctx context.Context, request mcp.CallToolRequest, dbSvc *vService.DMDBService) (*mcp.CallToolResult, error) {
dbName, _ := request.Params.Arguments["库名"].(string)
tableName, _ := request.Params.Arguments["表名"].(string)
kvStr, ok := request.Params.Arguments["键值对"].(string)
if !ok {
return mcp.NewToolResultError("缺少键值对参数"), nil
}
var kv map[string]interface{}
if err := json.Unmarshal([]byte(kvStr), &kv); err != nil {
return mcp.NewToolResultError("键值对必须是合法JSON: " + err.Error()), nil
}
if len(kv) == 0 {
return mcp.NewToolResultError("键值对不能为空"), nil
}
cols := make([]string, 0, len(kv))
placeholders := make([]string, 0, len(kv))
vals := make([]interface{}, 0, len(kv))
for k, v := range kv {
cols = append(cols, k)
placeholders = append(placeholders, "?")
vals = append(vals, v)
}
sql := "INSERT INTO " + dbName + "." + tableName + " (" + strings.Join(cols, ",") + ") VALUES (" + strings.Join(placeholders, ",") + ")"
n, err := dbSvc.ExecSQL(ctx, sql, vals...)
if err != nil {
return mcp.NewToolResultError(err.Error()), nil
}
return mcp.NewToolResultText("插入成功,影响行数: " + fmt.Sprintf("%d", n)), nil
}
// UpdateRows 更新记录
func UpdateRows(ctx context.Context, request mcp.CallToolRequest, dbSvc *vService.DMDBService) (*mcp.CallToolResult, error) {
dbName, _ := request.Params.Arguments["库名"].(string)
tableName, _ := request.Params.Arguments["表名"].(string)
setClause, _ := request.Params.Arguments["SET子句"].(string)
whereClause, _ := request.Params.Arguments["WHERE子句"].(string)
valsStr, ok := request.Params.Arguments["值列表"].(string)
if !ok {
return mcp.NewToolResultError("缺少值列表参数"), nil
}
var vals []interface{}
if err := json.Unmarshal([]byte(valsStr), &vals); err != nil {
return mcp.NewToolResultError("值列表必须是合法JSON数组: " + err.Error()), nil
}
sql := "UPDATE " + dbName + "." + tableName + " SET " + setClause + " WHERE " + whereClause
n, err := dbSvc.ExecSQL(ctx, sql, vals...)
if err != nil {
return mcp.NewToolResultError(err.Error()), nil
}
return mcp.NewToolResultText("更新成功,影响行数: " + fmt.Sprintf("%d", n)), nil
}
// DeleteRows 删除记录
func DeleteRows(ctx context.Context, request mcp.CallToolRequest, dbSvc *vService.DMDBService) (*mcp.CallToolResult, error) {
dbName, _ := request.Params.Arguments["库名"].(string)
tableName, _ := request.Params.Arguments["表名"].(string)
whereClause, _ := request.Params.Arguments["WHERE子句"].(string)
valsStr, ok := request.Params.Arguments["值列表"].(string)
if !ok {
return mcp.NewToolResultError("缺少值列表参数"), nil
}
var vals []interface{}
if err := json.Unmarshal([]byte(valsStr), &vals); err != nil {
return mcp.NewToolResultError("值列表必须是合法JSON数组: " + err.Error()), nil
}
sql := "DELETE FROM " + dbName + "." + tableName + " WHERE " + whereClause
n, err := dbSvc.ExecSQL(ctx, sql, vals...)
if err != nil {
return mcp.NewToolResultError(err.Error()), nil
}
return mcp.NewToolResultText("删除成功,影响行数: " + fmt.Sprintf("%d", n)), nil
}
// CreateTable 创建表
func CreateTable(ctx context.Context, request mcp.CallToolRequest, dbSvc *vService.DMDBService) (*mcp.CallToolResult, error) {
dbName, _ := request.Params.Arguments["库名"].(string)
tableName, _ := request.Params.Arguments["表名"].(string)
columnDefs, _ := request.Params.Arguments["字段定义"].(string)
if columnDefs == "" {
return mcp.NewToolResultError("缺少字段定义"), nil
}
if err := dbSvc.CreateTable(ctx, dbName, tableName, columnDefs); err != nil {
return mcp.NewToolResultError(err.Error()), nil
}
return mcp.NewToolResultText("创建表成功"), nil
}
// DropTable 删除表
func DropTable(ctx context.Context, request mcp.CallToolRequest, dbSvc *vService.DMDBService) (*mcp.CallToolResult, error) {
dbName, _ := request.Params.Arguments["库名"].(string)
tableName, _ := request.Params.Arguments["表名"].(string)
if err := dbSvc.DropTable(ctx, dbName, tableName); err != nil {
return mcp.NewToolResultError(err.Error()), nil
}
return mcp.NewToolResultText("删除表成功"), nil
}