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 }