commit 18e6bb2a3ba236511d1f1b0068bfa695ce8db445 Author: 张献维 Date: Mon Apr 28 12:48:51 2025 +0800 初始化 diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..115d712 --- /dev/null +++ b/.gitignore @@ -0,0 +1,3 @@ +/.idea/ +/main +/go.sum diff --git a/.goreleaser.yaml b/.goreleaser.yaml new file mode 100755 index 0000000..369fe7f --- /dev/null +++ b/.goreleaser.yaml @@ -0,0 +1,41 @@ +project_name: mcp-go-db +builds: + - env: + - CGO_ENABLED=0 + goos: + - linux + goarch: + - arm64 + id: "default" # 需要明确指定 id 供后续引用 + +nfpms: + - id: mcp-go-db + package_name: mcp-go-db + file_name_template: "{{ .ProjectName }}_{{ .Version }}_{{ .Os }}_{{ .Arch }}" + vendor: "Viviman" + homepage: "https://github.com/vivimanZhang/mcp-go-db" + maintainer: "vivib@foxmail.com" + description: "MCP-数据库小工具" + license: "MIT" + formats: + - deb + replaces: + - old-package-name + contents: + # 主程序安装到 /opt/mcp-go-db/mcp-go-db + - src: dist/default_linux_arm64/mcp-go-db + dst: /opt/apps/mcp-go-db/mcp-go-db + file_info: + mode: 0755 + # 配置文件安装到 /opt/apps/mcp-go-db/config.yaml + - src: config.yaml + dst: /opt/apps/mcp-go-db/config.yaml + type: config|noreplace # 标记为配置文件,升级时不会覆盖 + file_info: + mode: 0644 + + +release: + github: + owner: vivimanZhang + name: mcp-go-db \ No newline at end of file diff --git a/go.mod b/go.mod new file mode 100644 index 0000000..a71d769 --- /dev/null +++ b/go.mod @@ -0,0 +1,16 @@ +module mcp-go-db + +go 1.24 + +require ( + github.com/jmoiron/sqlx v1.4.0 + github.com/mark3labs/mcp-go v0.23.1 +) + +require ( + github.com/golang/snappy v1.0.0 // indirect + github.com/google/uuid v1.6.0 // indirect + github.com/spf13/cast v1.7.1 // indirect + github.com/yosida95/uritemplate/v3 v3.0.2 // indirect + golang.org/x/text v0.24.0 // indirect +) diff --git a/main.go b/main.go new file mode 100644 index 0000000..512e9f2 --- /dev/null +++ b/main.go @@ -0,0 +1,60 @@ +package main + +import ( + _ "dm" // 达梦数据库驱动 + "fmt" + "github.com/mark3labs/mcp-go/server" + "log" + "mcp-go-db/vConfig" + "mcp-go-db/vService" + "mcp-go-db/vTool" + "os" +) + +func main() { + // 开发 + // name := "dm://HY_ZGGL_FRAME:HY_ZGGL_FRAME@10.200.21.19:5296/DAMENG?schema=HY_ZGGL_FRAME" + + // 部署 + fmt.Println("程序名称:") + dsn := os.Args[1] + + fmt.Printf("启动数据库服务...\n") + // 加载配置 + config, err := vConfig.LoadConfig(dsn) + if err != nil { + log.Fatalf("配置文件加载失败: %v", err) + } + + // 初始化数据库连接 + dbSvc, err := vService.NewDMDBService(config.Database.DmDsn) + if err != nil { + fmt.Printf("数据库连接失败: %v\n", err) + return + } + defer func(dbSvc *vService.DMDBService) { + err := dbSvc.Close() + if err != nil { + log.Fatalf("发生异常: %v", err) + } + }(dbSvc) + fmt.Printf("数据库连接成功: %v\n", config.Database.DmDsn) + + // 创建MCP服务器 + s := server.NewMCPServer( + "数据库服务", + "1.0.0", + server.WithResourceCapabilities(true, true), + server.WithLogging(), + server.WithRecovery(), + ) + + // 添加数据库工具 + vTool.Run(s, dbSvc, config) + + // 启动MCP服务器 + if err := server.ServeStdio(s); err != nil { + fmt.Printf("服务器错误: %v\n", err) + } + fmt.Printf("服务已停止\n") +} diff --git a/vConfig/config.go b/vConfig/config.go new file mode 100644 index 0000000..9b96379 --- /dev/null +++ b/vConfig/config.go @@ -0,0 +1,75 @@ +package vConfig + +import ( + "fmt" + "net/url" + "strings" +) + +type Config struct { + Database struct { + DmDsn string `yaml:"dm_dsn"` + } `yaml:"database"` + User string + Pwd string + IP string + Port string + Schema string +} + +// LoadConfig 从指定路径加载 YAML 配置 +func LoadConfig(envPass string) (*Config, error) { + // 解析 YAML + var config Config + config.Database.DmDsn = envPass + + // 预处理:如果字符串以 dm:// 开头则替换为 http:// 便于 url.Parse 解析 + u, err := url.Parse(envPass) + if err != nil { + return nil, fmt.Errorf("解析连接字符串失败: %v", err) + } + + // 用户信息 + userInfo := u.User + if userInfo == nil { + return nil, fmt.Errorf("缺少用户信息") + } + config.User = userInfo.Username() + + // 获取密码 + if pwd, ok := userInfo.Password(); ok { + config.Pwd = pwd + } else { + return nil, fmt.Errorf("密码解析失败") + } + + // 拆分主机和端口 + host, port := u.Hostname(), u.Port() + if host == "" { + return nil, fmt.Errorf("缺少IP地址") + } + config.IP = host + + // 端口处理(默认端口5236) + if port == "" { + config.Port = "5236" // 达梦默认端口 + } else { + config.Port = port + } + + // 数据库名(路径部分) + if pathParts := strings.Split(strings.Trim(u.Path, "/"), "/"); len(pathParts) > 0 { + config.Schema = pathParts[0] + } else { + return nil, fmt.Errorf("缺少数据库名") + } + + // 获取schema参数 + schema := u.Query().Get("schema") + if schema == "" { + return nil, fmt.Errorf("缺少schema参数") + } + config.Schema = schema + + return &config, nil +} diff --git a/vConfig/config_test.go b/vConfig/config_test.go new file mode 100644 index 0000000..9b9363b --- /dev/null +++ b/vConfig/config_test.go @@ -0,0 +1,93 @@ +package vConfig + +import ( + "testing" +) + +func TestParseDMConnectionString(t *testing.T) { + tests := []struct { + name string + connStr string + want *Config + wantErr bool + errMsg string + }{ + { + name: "normal case with default port", + connStr: "dm://user:pass@127.0.0.1/DAMENG?schema=MY_SCHEMA", + want: &Config{ + User: "user", + Pwd: "pass", + IP: "127.0.0.1", + Port: "5236", + Schema: "MY_SCHEMA", + }, + wantErr: false, + }, + { + name: "explicit port", + connStr: "dm://user:pass@10.200.21.19:5296/DAMENG?schema=HY_ZGGL_GZ", + want: &Config{ + User: "user", + Pwd: "pass", + IP: "10.200.21.19", + Port: "5296", + Schema: "HY_ZGGL_GZ", + }, + wantErr: false, + }, + { + name: "missing schema", + connStr: "dm://user:pass@127.0.0.1/DAMENG", + wantErr: true, + errMsg: "缺少schema参数", + }, + { + name: "invalid format", + connStr: "dm://invalid_user_info_format", + wantErr: true, + errMsg: "缺少用户信息", + }, + { + name: "empty password", + connStr: "dm://user@127.0.0.1/DAMENG?schema=TEST", + wantErr: true, + errMsg: "密码解析失败", + }, + { + name: "special characters in password", + connStr: "dm://user:p%40ssw0rd%21%40%23@127.0.0.1/DAMENG?schema=TEST", + want: &Config{ + User: "user", + Pwd: "p@ssw0rd!@#", + IP: "127.0.0.1", + Port: "5236", + Schema: "TEST", + }, + wantErr: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got, err := LoadConfig(tt.connStr) + if (err != nil) != tt.wantErr { + t.Errorf("ParseDMConnectionString() error = %v, wantErr %v", err, tt.wantErr) + return + } + if tt.wantErr { + if err.Error() != tt.errMsg { + t.Errorf("ParseDMConnectionString() error = %v, wantErrMsg %v", err.Error(), tt.errMsg) + } + return + } + if got.User != tt.want.User || + got.Pwd != tt.want.Pwd || + got.IP != tt.want.IP || + got.Port != tt.want.Port || + got.Schema != tt.want.Schema { + t.Errorf("ParseDMConnectionString() = %+v, want %+v", got, tt.want) + } + }) + } +} diff --git a/vService/DmService.go b/vService/DmService.go new file mode 100644 index 0000000..4accdcd --- /dev/null +++ b/vService/DmService.go @@ -0,0 +1,101 @@ +package vService + +import ( + "context" + _ "dm" // 达梦数据库驱动 + "fmt" + "github.com/jmoiron/sqlx" + "log" + "strings" +) + +// DMDBService 数据库服务实现 +type DMDBService struct { + db *sqlx.DB +} + +func NewDMDBService(dsn string) (*DMDBService, error) { + db, err := sqlx.Connect("dm", dsn) + if err != nil { + return nil, err + } + return &DMDBService{db: db}, nil +} + +func (s *DMDBService) Close() error { + return s.db.Close() +} + +func (s *DMDBService) FetchData(ctx context.Context, dbName, tableName string, limit int) ([]map[string]interface{}, error) { + query := fmt.Sprintf("SELECT * FROM %s.%s WHERE ROWNUM <= %d", dbName, tableName, limit) + + results, i, err2 := s.queryBySQL(ctx, query) + if err2 != nil { + return i, err2 + } + return results, nil +} + +func (s *DMDBService) queryBySQL(ctx context.Context, query string) ([]map[string]interface{}, []map[string]interface{}, error) { + rows, err := s.db.QueryxContext(ctx, query) + if err != nil { + return nil, nil, err + } + defer rows.Close() + + var results []map[string]interface{} + for rows.Next() { + row := make(map[string]interface{}) + if err := rows.MapScan(row); err != nil { + return nil, nil, err + } + results = append(results, row) + } + return results, nil, nil +} + +func (s *DMDBService) AddColumn(ctx context.Context, dbName, tableName, columnName, columnType string) error { + query := fmt.Sprintf("ALTER TABLE %s.%s ADD %s %s", dbName, tableName, columnName, columnType) + if _, err := s.db.ExecContext(ctx, query); err != nil { + return err + } + + return nil +} + +func (s *DMDBService) AddComment(ctx context.Context, dbName, tableName, columnName, comment string) error { + if comment != "" { + commentQuery := fmt.Sprintf("COMMENT ON COLUMN %s.%s.%s IS '%s'", + dbName, tableName, columnName, strings.ReplaceAll(comment, "'", "''")) + if _, err := s.db.ExecContext(ctx, commentQuery); err != nil { + return err + } + } + return nil +} + +func (s *DMDBService) DropColumn(ctx context.Context, dbName, tableName, columnName string) error { + query := fmt.Sprintf("ALTER TABLE %s.%s DROP COLUMN %s", dbName, tableName, columnName) + _, err := s.db.ExecContext(ctx, query) + return err +} + +func (s *DMDBService) AllTable(ctx context.Context, dbName string) (interface{}, error) { + query := fmt.Sprintf("SELECT TABLE_NAME FROM ALL_TABLES WHERE OWNER = '%s'", dbName) + rows, err := s.db.QueryxContext(ctx, query) + if err != nil { + return nil, err + } + defer func(rows *sqlx.Rows) { + err := rows.Close() + if err != nil { + log.Fatalf("发生异常: %v", err) + } + }(rows) + + results, i, err := s.queryBySQL(ctx, query) + if err != nil { + return i, err + } + return results, nil +} diff --git a/vTool/dbTool.go b/vTool/dbTool.go new file mode 100644 index 0000000..c192749 --- /dev/null +++ b/vTool/dbTool.go @@ -0,0 +1,115 @@ +package vTool + +import ( + "context" + "encoding/json" + "github.com/mark3labs/mcp-go/mcp" + "mcp-go-db/vConfig" + "mcp-go-db/vService" +) + +// 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 +} + +// 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 +} diff --git a/vTool/mcpTools.go b/vTool/mcpTools.go new file mode 100644 index 0000000..9f8e49c --- /dev/null +++ b/vTool/mcpTools.go @@ -0,0 +1,140 @@ +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), + ), + ) +} diff --git a/说明.md b/说明.md new file mode 100644 index 0000000..1b6d108 --- /dev/null +++ b/说明.md @@ -0,0 +1,23 @@ +# 中文版本(MCP-GO Demo) + + +配置文件 + +```json +{ + "mcpServers": { + "calculator": { + "name": "计算小样", + "version": "1.0.0", + "command": "~/Music/go-code/mcp-go-demo/main", + "args": [], + "description": "提供加减乘除运算服务", + "enabled": true, + "autoApprove": [ + "计算" + ] + } + } +} + +``` \ No newline at end of file