后端重构

This commit is contained in:
Yvan
2026-01-14 14:34:49 +08:00
parent f5f77b32fe
commit 30842b42e5
274 changed files with 39754 additions and 1005 deletions
+41
View File
@@ -0,0 +1,41 @@
package client
import (
"context"
"errors"
mcpClient "github.com/mark3labs/mcp-go/client"
"github.com/mark3labs/mcp-go/mcp"
)
// NewClient 创建MCP客户端
func NewClient(baseUrl, name, version, serverName string) (*mcpClient.Client, error) {
client, err := mcpClient.NewSSEMCPClient(baseUrl)
if err != nil {
return nil, err
}
ctx := context.Background()
// 启动client
if err := client.Start(ctx); err != nil {
return nil, err
}
// 初始化
initRequest := mcp.InitializeRequest{}
initRequest.Params.ProtocolVersion = mcp.LATEST_PROTOCOL_VERSION
initRequest.Params.ClientInfo = mcp.Implementation{
Name: name,
Version: version,
}
result, err := client.Initialize(ctx, initRequest)
if err != nil {
return nil, err
}
if result.ServerInfo.Name != serverName {
return nil, errors.New("server name mismatch")
}
return client, nil
}
+205
View File
@@ -0,0 +1,205 @@
package mcpTool
import (
"context"
"encoding/json"
"errors"
"fmt"
"github.com/mark3labs/mcp-go/mcp"
"go.uber.org/zap"
"gorm.io/gorm"
"kra/internal/data/model"
)
// 注册工具
func init() {
RegisterTool(&ApiCreator{})
}
// ApiCreateRequest API创建请求结构
type ApiCreateRequest struct {
Path string `json:"path"` // API路径
Description string `json:"description"` // API中文描述
ApiGroup string `json:"apiGroup"` // API组
Method string `json:"method"` // HTTP方法
}
// ApiCreateResponse API创建响应结构
type ApiCreateResponse struct {
Success bool `json:"success"`
Message string `json:"message"`
ApiID uint `json:"apiId"`
Path string `json:"path"`
Method string `json:"method"`
}
// ApiCreator API创建工具
type ApiCreator struct {
db *gorm.DB
logger *zap.Logger
}
// SetDependencies 设置依赖
func (a *ApiCreator) SetDependencies(db *gorm.DB, logger *zap.Logger) {
a.db = db
a.logger = logger
}
// New 创建API创建工具
func (a *ApiCreator) New() mcp.Tool {
return mcp.NewTool("create_api",
mcp.WithDescription(`创建后端API记录,用于AI编辑器自动添加API接口时自动创建对应的API权限记录。
**重要限制:**
- 当使用kra_auto_generate工具且needCreatedModules=true时,模块创建会自动生成API权限,不应调用此工具
- 仅在以下情况使用:1) 单独创建API(不涉及模块创建);2) AI编辑器自动添加API;3) router下的文件产生路径变化时`),
mcp.WithString("path",
mcp.Required(),
mcp.Description("API路径,如:/user/create"),
),
mcp.WithString("description",
mcp.Required(),
mcp.Description("API中文描述,如:创建用户"),
),
mcp.WithString("apiGroup",
mcp.Required(),
mcp.Description("API组名称,用于分类管理,如:用户管理"),
),
mcp.WithString("method",
mcp.Description("HTTP方法"),
mcp.DefaultString("POST"),
),
mcp.WithString("apis",
mcp.Description("批量创建API的JSON字符串,格式:[{\"path\":\"/user/create\",\"description\":\"创建用户\",\"apiGroup\":\"用户管理\",\"method\":\"POST\"}]"),
),
)
}
// Handle 处理API创建请求
func (a *ApiCreator) Handle(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
if a.db == nil {
return nil, errors.New("数据库连接未初始化")
}
args := request.GetArguments()
var apis []ApiCreateRequest
// 检查是否是批量创建
if apisStr, ok := args["apis"].(string); ok && apisStr != "" {
if err := json.Unmarshal([]byte(apisStr), &apis); err != nil {
return nil, fmt.Errorf("apis 参数格式错误: %v", err)
}
} else {
// 单个API创建
path, ok := args["path"].(string)
if !ok || path == "" {
return nil, errors.New("path 参数是必需的")
}
description, ok := args["description"].(string)
if !ok || description == "" {
return nil, errors.New("description 参数是必需的")
}
apiGroup, ok := args["apiGroup"].(string)
if !ok || apiGroup == "" {
return nil, errors.New("apiGroup 参数是必需的")
}
method := "POST"
if val, ok := args["method"].(string); ok && val != "" {
method = val
}
apis = append(apis, ApiCreateRequest{
Path: path,
Description: description,
ApiGroup: apiGroup,
Method: method,
})
}
if len(apis) == 0 {
return nil, errors.New("没有要创建的API")
}
// 创建API记录
var responses []ApiCreateResponse
successCount := 0
for _, apiReq := range apis {
api := model.SysAPI{
Path: apiReq.Path,
Description: apiReq.Description,
APIGroup: apiReq.ApiGroup,
Method: apiReq.Method,
}
err := a.db.Create(&api).Error
if err != nil {
if a.logger != nil {
a.logger.Warn("创建API失败",
zap.String("path", apiReq.Path),
zap.String("method", apiReq.Method),
zap.Error(err))
}
responses = append(responses, ApiCreateResponse{
Success: false,
Message: fmt.Sprintf("创建API失败: %v", err),
Path: apiReq.Path,
Method: apiReq.Method,
})
} else {
// 获取创建的API ID
var createdApi model.SysAPI
err = a.db.Where("path = ? AND method = ?", apiReq.Path, apiReq.Method).First(&createdApi).Error
if err != nil && a.logger != nil {
a.logger.Warn("获取创建的API ID失败", zap.Error(err))
}
responses = append(responses, ApiCreateResponse{
Success: true,
Message: fmt.Sprintf("成功创建API %s %s", apiReq.Method, apiReq.Path),
ApiID: uint(createdApi.ID),
Path: apiReq.Path,
Method: apiReq.Method,
})
successCount++
}
}
// 构建总体响应
var resultMessage string
if len(apis) == 1 {
resultMessage = responses[0].Message
} else {
resultMessage = fmt.Sprintf("批量创建API完成,成功 %d 个,失败 %d 个", successCount, len(apis)-successCount)
}
result := map[string]interface{}{
"success": successCount > 0,
"message": resultMessage,
"totalCount": len(apis),
"successCount": successCount,
"failedCount": len(apis) - successCount,
"details": responses,
}
resultJSON, err := json.MarshalIndent(result, "", " ")
if err != nil {
return nil, fmt.Errorf("序列化结果失败: %v", err)
}
return &mcp.CallToolResult{
Content: []mcp.Content{
mcp.TextContent{
Type: "text",
Text: fmt.Sprintf("API创建结果:\n\n%s", string(resultJSON)),
},
},
}, nil
}
+169
View File
@@ -0,0 +1,169 @@
package mcpTool
import (
"context"
"encoding/json"
"fmt"
"github.com/gin-gonic/gin"
"github.com/mark3labs/mcp-go/mcp"
"go.uber.org/zap"
"gorm.io/gorm"
"kra/internal/data/model"
)
// 注册工具
func init() {
RegisterTool(&ApiLister{})
}
// ApiInfo API信息结构
type ApiInfo struct {
ID uint `json:"id,omitempty"` // 数据库ID(仅数据库API有)
Path string `json:"path"` // API路径
Description string `json:"description,omitempty"` // API描述
ApiGroup string `json:"apiGroup,omitempty"` // API组
Method string `json:"method"` // HTTP方法
Source string `json:"source"` // 来源:database 或 gin
}
// ApiListResponse API列表响应结构
type ApiListResponse struct {
Success bool `json:"success"`
Message string `json:"message"`
DatabaseApis []ApiInfo `json:"databaseApis"` // 数据库中的API
GinApis []ApiInfo `json:"ginApis"` // gin框架中的API
TotalCount int `json:"totalCount"` // 总数量
}
// ApiLister API列表工具
type ApiLister struct {
db *gorm.DB
logger *zap.Logger
routers gin.RoutesInfo
}
// SetDependencies 设置依赖
func (a *ApiLister) SetDependencies(db *gorm.DB, logger *zap.Logger, routers gin.RoutesInfo) {
a.db = db
a.logger = logger
a.routers = routers
}
// New 创建API列表工具
func (a *ApiLister) New() mcp.Tool {
return mcp.NewTool("list_all_apis",
mcp.WithDescription(`获取系统中所有的API接口,分为两组:
**功能说明:**
- 返回数据库中已注册的API列表
- 返回gin框架中实际注册的路由API列表
- 帮助前端判断是使用现有API还是需要创建新的API,如果api在前端未使用且需要前端调用的时候,请到api文件夹下对应模块的js中添加方法并暴露给当前业务调用
**返回数据结构:**
- databaseApis: 数据库中的API记录(包含ID、描述、分组等完整信息)
- ginApis: gin路由中的API(仅包含路径和方法),需要AI根据路径自行揣摩路径的业务含义,例如:/api/user/:id 表示根据用户ID获取用户信息`),
mcp.WithString("_placeholder",
mcp.Description("占位符,防止json schema校验失败"),
),
)
}
// Handle 处理API列表请求
func (a *ApiLister) Handle(_ context.Context, _ mcp.CallToolRequest) (*mcp.CallToolResult, error) {
// 获取数据库中的API
databaseApis, err := a.getDatabaseApis()
if err != nil {
if a.logger != nil {
a.logger.Error("获取数据库API失败", zap.Error(err))
}
errorResponse := ApiListResponse{
Success: false,
Message: "获取数据库API失败: " + err.Error(),
}
resultJSON, _ := json.Marshal(errorResponse)
return &mcp.CallToolResult{
Content: []mcp.Content{
mcp.TextContent{
Type: "text",
Text: string(resultJSON),
},
},
}, nil
}
// 获取gin路由中的API
ginApis := a.getGinApis()
// 构建响应
response := ApiListResponse{
Success: true,
Message: "获取API列表成功",
DatabaseApis: databaseApis,
GinApis: ginApis,
TotalCount: len(databaseApis) + len(ginApis),
}
if a.logger != nil {
a.logger.Info("API列表获取成功",
zap.Int("数据库API数量", len(databaseApis)),
zap.Int("gin路由API数量", len(ginApis)),
zap.Int("总数量", response.TotalCount))
}
resultJSON, err := json.Marshal(response)
if err != nil {
return nil, fmt.Errorf("序列化结果失败: %v", err)
}
return &mcp.CallToolResult{
Content: []mcp.Content{
mcp.TextContent{
Type: "text",
Text: string(resultJSON),
},
},
}, nil
}
// getDatabaseApis 获取数据库中的所有API
func (a *ApiLister) getDatabaseApis() ([]ApiInfo, error) {
if a.db == nil {
return nil, fmt.Errorf("数据库连接未初始化")
}
var apis []model.SysAPI
err := a.db.Model(&model.SysAPI{}).Order("api_group ASC, path ASC").Find(&apis).Error
if err != nil {
return nil, err
}
// 转换为ApiInfo格式
var result []ApiInfo
for _, api := range apis {
result = append(result, ApiInfo{
ID: uint(api.ID),
Path: api.Path,
Description: api.Description,
ApiGroup: api.APIGroup,
Method: api.Method,
Source: "database",
})
}
return result, nil
}
// getGinApis 获取gin路由中的所有API
func (a *ApiLister) getGinApis() []ApiInfo {
var result []ApiInfo
for _, route := range a.routers {
result = append(result, ApiInfo{
Path: route.Path,
Method: route.Method,
Source: "gin",
})
}
return result
}
+240
View File
@@ -0,0 +1,240 @@
package mcpTool
import (
"context"
"encoding/json"
"errors"
"fmt"
"github.com/mark3labs/mcp-go/mcp"
"go.uber.org/zap"
"gorm.io/gorm"
"kra/internal/data/model"
)
func init() {
RegisterTool(&DictionaryOptionsGenerator{})
}
// DictionaryOptionsGenerator 字典选项生成器
type DictionaryOptionsGenerator struct {
db *gorm.DB
logger *zap.Logger
}
// SetDependencies 设置依赖
func (d *DictionaryOptionsGenerator) SetDependencies(db *gorm.DB, logger *zap.Logger) {
d.db = db
d.logger = logger
}
// DictionaryOption 字典选项结构
type DictionaryOption struct {
Label string `json:"label"`
Value string `json:"value"`
Sort int `json:"sort"`
}
// DictionaryGenerateRequest 字典生成请求
type DictionaryGenerateRequest struct {
DictType string `json:"dictType"` // 字典类型
FieldDesc string `json:"fieldDesc"` // 字段描述
Options []DictionaryOption `json:"options"` // AI生成的字典选项
DictName string `json:"dictName"` // 字典名称(可选)
Description string `json:"description"` // 字典描述(可选)
}
// DictionaryGenerateResponse 字典生成响应
type DictionaryGenerateResponse struct {
Success bool `json:"success"`
Message string `json:"message"`
DictType string `json:"dictType"`
OptionsCount int `json:"optionsCount"`
}
// New 返回工具注册信息
func (d *DictionaryOptionsGenerator) New() mcp.Tool {
return mcp.NewTool("generate_dictionary_options",
mcp.WithDescription("智能生成字典选项并自动创建字典和字典详情"),
mcp.WithString("dictType",
mcp.Required(),
mcp.Description("字典类型,用于标识字典的唯一性"),
),
mcp.WithString("fieldDesc",
mcp.Required(),
mcp.Description("字段描述,用于AI理解字段含义"),
),
mcp.WithString("options",
mcp.Required(),
mcp.Description("字典选项JSON字符串,格式:[{\"label\":\"显示名\",\"value\":\"值\",\"sort\":1}]"),
),
mcp.WithString("dictName",
mcp.Description("字典名称,如果不提供将自动生成"),
),
mcp.WithString("description",
mcp.Description("字典描述"),
),
)
}
// Handle 处理工具调用
func (d *DictionaryOptionsGenerator) Handle(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
if d.db == nil {
return nil, errors.New("数据库连接未初始化")
}
// 解析请求参数
args := request.GetArguments()
dictType, ok := args["dictType"].(string)
if !ok || dictType == "" {
return nil, errors.New("dictType 参数是必需的")
}
fieldDesc, ok := args["fieldDesc"].(string)
if !ok || fieldDesc == "" {
return nil, errors.New("fieldDesc 参数是必需的")
}
optionsStr, ok := args["options"].(string)
if !ok || optionsStr == "" {
return nil, errors.New("options 参数是必需的")
}
// 解析options JSON字符串
var options []DictionaryOption
if err := json.Unmarshal([]byte(optionsStr), &options); err != nil {
return nil, fmt.Errorf("options 参数格式错误: %v", err)
}
if len(options) == 0 {
return nil, errors.New("options 不能为空")
}
dictName, _ := args["dictName"].(string)
description, _ := args["description"].(string)
// 构建请求对象
req := &DictionaryGenerateRequest{
DictType: dictType,
FieldDesc: fieldDesc,
Options: options,
DictName: dictName,
Description: description,
}
// 创建字典
response, err := d.createDictionaryWithOptions(ctx, req)
if err != nil {
return nil, err
}
// 构建响应
resultJSON, err := json.MarshalIndent(response, "", " ")
if err != nil {
return nil, fmt.Errorf("序列化结果失败: %v", err)
}
return &mcp.CallToolResult{
Content: []mcp.Content{
mcp.TextContent{
Type: "text",
Text: fmt.Sprintf("字典选项生成结果:\n\n%s", string(resultJSON)),
},
},
}, nil
}
// createDictionaryWithOptions 创建字典和字典选项
func (d *DictionaryOptionsGenerator) createDictionaryWithOptions(ctx context.Context, req *DictionaryGenerateRequest) (*DictionaryGenerateResponse, error) {
// 检查字典是否已存在
exists, err := d.checkDictionaryExists(req.DictType)
if err != nil {
return nil, fmt.Errorf("检查字典是否存在失败: %v", err)
}
if exists {
return &DictionaryGenerateResponse{
Success: false,
Message: fmt.Sprintf("字典 %s 已存在,跳过创建", req.DictType),
DictType: req.DictType,
OptionsCount: 0,
}, nil
}
// 生成字典名称
dictName := req.DictName
if dictName == "" {
dictName = d.generateDictionaryName(req.DictType, req.FieldDesc)
}
// 创建字典
dictionary := model.SysDictionary{
Name: dictName,
Type: req.DictType,
Status: true,
Desc: req.Description,
}
err = d.db.Create(&dictionary).Error
if err != nil {
return nil, fmt.Errorf("创建字典失败: %v", err)
}
// 获取刚创建的字典ID
var createdDict model.SysDictionary
err = d.db.Where("type = ?", req.DictType).First(&createdDict).Error
if err != nil {
return nil, fmt.Errorf("获取创建的字典失败: %v", err)
}
// 创建字典详情项
successCount := 0
for _, option := range req.Options {
dictionaryDetail := model.SysDictionaryDetail{
Label: option.Label,
Value: option.Value,
Status: true,
Sort: int64(option.Sort),
SysDictionaryID: createdDict.ID,
}
err = d.db.Create(&dictionaryDetail).Error
if err != nil {
if d.logger != nil {
d.logger.Warn("创建字典详情项失败", zap.Error(err))
}
} else {
successCount++
}
}
return &DictionaryGenerateResponse{
Success: true,
Message: fmt.Sprintf("成功创建字典 %s,包含 %d 个选项", req.DictType, successCount),
DictType: req.DictType,
OptionsCount: successCount,
}, nil
}
// checkDictionaryExists 检查字典是否存在
func (d *DictionaryOptionsGenerator) checkDictionaryExists(dictType string) (bool, error) {
var dictionary model.SysDictionary
err := d.db.Where("type = ?", dictType).First(&dictionary).Error
if err != nil {
if errors.Is(err, gorm.ErrRecordNotFound) {
return false, nil // 字典不存在
}
return false, err // 其他错误
}
return true, nil // 字典存在
}
// generateDictionaryName 生成字典名称
func (d *DictionaryOptionsGenerator) generateDictionaryName(dictType, fieldDesc string) string {
if fieldDesc != "" {
return fmt.Sprintf("%s字典", fieldDesc)
}
return fmt.Sprintf("%s字典", dictType)
}
+262
View File
@@ -0,0 +1,262 @@
package mcpTool
import (
"context"
"encoding/json"
"fmt"
"github.com/mark3labs/mcp-go/mcp"
"go.uber.org/zap"
"gorm.io/gorm"
"kra/internal/data/model"
)
// 注册工具
func init() {
RegisterTool(&DictionaryQuery{})
}
// DictionaryPre 字典预览结构
type DictionaryPre struct {
Type string `json:"type"` // 字典名(英)
Desc string `json:"desc"` // 描述
}
// DictionaryInfo 字典信息结构
type DictionaryInfo struct {
ID uint `json:"id"`
Name string `json:"name"` // 字典名(中)
Type string `json:"type"` // 字典名(英)
Status bool `json:"status"` // 状态
Desc string `json:"desc"` // 描述
Details []DictionaryDetailInfo `json:"details"` // 字典详情
}
// DictionaryDetailInfo 字典详情信息结构
type DictionaryDetailInfo struct {
ID uint `json:"id"`
Label string `json:"label"` // 展示值
Value string `json:"value"` // 字典值
Extend string `json:"extend"` // 扩展值
Status bool `json:"status"` // 启用状态
Sort int64 `json:"sort"` // 排序标记
}
// DictionaryQueryResponse 字典查询响应结构
type DictionaryQueryResponse struct {
Success bool `json:"success"`
Message string `json:"message"`
Total int `json:"total"`
Dictionaries []DictionaryInfo `json:"dictionaries"`
}
// DictionaryQuery 字典查询工具
type DictionaryQuery struct {
db *gorm.DB
logger *zap.Logger
}
// SetDependencies 设置依赖
func (d *DictionaryQuery) SetDependencies(db *gorm.DB, logger *zap.Logger) {
d.db = db
d.logger = logger
}
// New 创建字典查询工具
func (d *DictionaryQuery) New() mcp.Tool {
return mcp.NewTool("query_dictionaries",
mcp.WithDescription("查询系统中所有的字典和字典属性,用于AI生成逻辑时了解可用的字典选项"),
mcp.WithString("dictType",
mcp.Description("可选:指定字典类型进行精确查询,如果不提供则返回所有字典"),
),
mcp.WithBoolean("includeDisabled",
mcp.Description("是否包含已禁用的字典和字典项,默认为false(只返回启用的)"),
),
mcp.WithBoolean("detailsOnly",
mcp.Description("是否只返回字典详情信息(不包含字典基本信息),默认为false"),
),
)
}
// Handle 处理字典查询请求
func (d *DictionaryQuery) Handle(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
if d.db == nil {
return &mcp.CallToolResult{
Content: []mcp.Content{
mcp.NewTextContent(`{"success": false, "message": "数据库连接未初始化", "total": 0, "dictionaries": []}`),
},
}, nil
}
args := request.GetArguments()
// 获取参数
dictType := ""
if val, ok := args["dictType"].(string); ok {
dictType = val
}
includeDisabled := false
if val, ok := args["includeDisabled"].(bool); ok {
includeDisabled = val
}
detailsOnly := false
if val, ok := args["detailsOnly"].(bool); ok {
detailsOnly = val
}
var dictionaries []DictionaryInfo
if dictType != "" {
// 查询指定类型的字典
var sysDictionary model.SysDictionary
db := d.db.Model(&model.SysDictionary{}).Where("type = ?", dictType)
if !includeDisabled {
db = db.Where("status = ?", true)
}
err := db.First(&sysDictionary).Error
if err != nil {
if d.logger != nil {
d.logger.Error("查询字典失败", zap.Error(err))
}
return &mcp.CallToolResult{
Content: []mcp.Content{
mcp.NewTextContent(fmt.Sprintf(`{"success": false, "message": "查询字典失败: %v", "total": 0, "dictionaries": []}`, err.Error())),
},
}, nil
}
// 转换为响应格式
dictInfo := DictionaryInfo{
ID: uint(sysDictionary.ID),
Name: sysDictionary.Name,
Type: sysDictionary.Type,
Status: sysDictionary.Status,
Desc: sysDictionary.Desc,
}
// 查询字典详情
var details []model.SysDictionaryDetail
detailDB := d.db.Model(&model.SysDictionaryDetail{}).Where("sys_dictionary_id = ?", sysDictionary.ID)
if !includeDisabled {
detailDB = detailDB.Where("status = ?", true)
}
detailDB.Order("sort").Find(&details)
for _, detail := range details {
dictInfo.Details = append(dictInfo.Details, DictionaryDetailInfo{
ID: uint(detail.ID),
Label: detail.Label,
Value: detail.Value,
Extend: detail.Extend,
Status: detail.Status,
Sort: detail.Sort,
})
}
dictionaries = append(dictionaries, dictInfo)
} else {
// 查询所有字典
var sysDictionaries []model.SysDictionary
db := d.db.Model(&model.SysDictionary{})
if !includeDisabled {
db = db.Where("status = ?", true)
}
err := db.Find(&sysDictionaries).Error
if err != nil {
if d.logger != nil {
d.logger.Error("查询字典列表失败", zap.Error(err))
}
return &mcp.CallToolResult{
Content: []mcp.Content{
mcp.NewTextContent(fmt.Sprintf(`{"success": false, "message": "查询字典列表失败: %v", "total": 0, "dictionaries": []}`, err.Error())),
},
}, nil
}
// 转换为响应格式
for _, dict := range sysDictionaries {
dictInfo := DictionaryInfo{
ID: uint(dict.ID),
Name: dict.Name,
Type: dict.Type,
Status: dict.Status,
Desc: dict.Desc,
}
// 查询字典详情
var details []model.SysDictionaryDetail
detailDB := d.db.Model(&model.SysDictionaryDetail{}).Where("sys_dictionary_id = ?", dict.ID)
if !includeDisabled {
detailDB = detailDB.Where("status = ?", true)
}
detailDB.Order("sort").Find(&details)
for _, detail := range details {
dictInfo.Details = append(dictInfo.Details, DictionaryDetailInfo{
ID: uint(detail.ID),
Label: detail.Label,
Value: detail.Value,
Extend: detail.Extend,
Status: detail.Status,
Sort: detail.Sort,
})
}
dictionaries = append(dictionaries, dictInfo)
}
}
// 如果只需要详情信息,则提取所有详情
if detailsOnly {
var allDetails []DictionaryDetailInfo
for _, dict := range dictionaries {
allDetails = append(allDetails, dict.Details...)
}
response := map[string]interface{}{
"success": true,
"message": "查询字典详情成功",
"total": len(allDetails),
"details": allDetails,
}
responseJSON, _ := json.Marshal(response)
return &mcp.CallToolResult{
Content: []mcp.Content{
mcp.NewTextContent(string(responseJSON)),
},
}, nil
}
// 构建响应
response := DictionaryQueryResponse{
Success: true,
Message: "查询字典成功",
Total: len(dictionaries),
Dictionaries: dictionaries,
}
responseJSON, err := json.Marshal(response)
if err != nil {
if d.logger != nil {
d.logger.Error("序列化响应失败", zap.Error(err))
}
return &mcp.CallToolResult{
Content: []mcp.Content{
mcp.NewTextContent(fmt.Sprintf(`{"success": false, "message": "序列化响应失败: %v", "total": 0, "dictionaries": []}`, err.Error())),
},
}, nil
}
return &mcp.CallToolResult{
Content: []mcp.Content{
mcp.NewTextContent(string(responseJSON)),
},
}, nil
}
+63
View File
@@ -0,0 +1,63 @@
package mcpTool
import (
"context"
"github.com/mark3labs/mcp-go/mcp"
"github.com/mark3labs/mcp-go/server"
"go.uber.org/zap"
"gorm.io/gorm"
)
// McpTool 定义了MCP工具必须实现的接口
type McpTool interface {
// Handle 返回工具调用信息
Handle(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error)
// New 返回工具注册信息
New() mcp.Tool
}
// DependencyInjectable 定义了支持依赖注入的工具接口
type DependencyInjectable interface {
SetDependencies(db *gorm.DB, logger *zap.Logger)
}
// ConfigInjectable 定义了支持配置注入的工具接口
type ConfigInjectable interface {
SetDependencies(db *gorm.DB, logger *zap.Logger, config *AutoCodeConfig)
}
// 工具注册表
var toolRegister = make(map[string]McpTool)
// RegisterTool 供工具在init时调用,将自己注册到工具注册表中
func RegisterTool(tool McpTool) {
mcpTool := tool.New()
toolRegister[mcpTool.Name] = tool
}
// RegisterAllTools 将所有注册的工具注册到MCP服务中
func RegisterAllTools(mcpServer *server.MCPServer) {
for _, tool := range toolRegister {
mcpServer.AddTool(tool.New(), tool.Handle)
}
}
// RegisterAllToolsWithDependencies 将所有注册的工具注册到MCP服务中,并注入依赖
func RegisterAllToolsWithDependencies(mcpServer *server.MCPServer, db *gorm.DB, logger *zap.Logger, config *AutoCodeConfig) {
for _, tool := range toolRegister {
// 尝试注入配置依赖
if injectable, ok := tool.(ConfigInjectable); ok {
injectable.SetDependencies(db, logger, config)
} else if injectable, ok := tool.(DependencyInjectable); ok {
// 尝试注入基本依赖
injectable.SetDependencies(db, logger)
}
mcpServer.AddTool(tool.New(), tool.Handle)
}
}
// GetRegisteredTools 获取所有已注册的工具
func GetRegisteredTools() map[string]McpTool {
return toolRegister
}
+552
View File
@@ -0,0 +1,552 @@
package mcpTool
import (
"context"
"encoding/json"
"errors"
"fmt"
"os"
"path/filepath"
"strings"
"github.com/mark3labs/mcp-go/mcp"
"go.uber.org/zap"
"gorm.io/gorm"
"kra/internal/data/model"
)
// 注册工具
func init() {
RegisterTool(&KRAAnalyzer{})
}
// KRAAnalyzer KRA分析器 - 用于分析当前功能是否需要创建独立的package和module
type KRAAnalyzer struct {
db *gorm.DB
logger *zap.Logger
config *AutoCodeConfig
}
// AutoCodeConfig 自动代码生成配置
type AutoCodeConfig struct {
Root string
Server string
Web string
Module string
}
// SetDependencies 设置依赖
func (g *KRAAnalyzer) SetDependencies(db *gorm.DB, logger *zap.Logger, config *AutoCodeConfig) {
g.db = db
g.logger = logger
g.config = config
}
// AnalyzeRequest 分析请求结构体
type AnalyzeRequest struct {
Requirement string `json:"requirement" binding:"required"` // 用户需求描述
}
// AnalyzeResponse 分析响应结构体
type AnalyzeResponse struct {
ExistingPackages []PackageInfo `json:"existingPackages"` // 现有包信息
PredesignedModules []PredesignedModuleInfo `json:"predesignedModules"` // 预设计模块信息
Dictionaries []DictionaryPre `json:"dictionaries"` // 字典信息
CleanupInfo *CleanupInfo `json:"cleanupInfo"` // 清理信息(如果有)
}
// ModuleInfo 模块信息
type ModuleInfo struct {
ModuleName string `json:"moduleName"` // 模块名称
PackageName string `json:"packageName"` // 包名
Template string `json:"template"` // 模板类型
StructName string `json:"structName"` // 结构体名称
TableName string `json:"tableName"` // 表名
Description string `json:"description"` // 描述
FilePaths []string `json:"filePaths"` // 相关文件路径
}
// PackageInfo 包信息
type PackageInfo struct {
PackageName string `json:"packageName"` // 包名
Template string `json:"template"` // 模板类型
Label string `json:"label"` // 标签
Desc string `json:"desc"` // 描述
Module string `json:"module"` // 模块
IsEmpty bool `json:"isEmpty"` // 是否为空包
}
// PredesignedModuleInfo 预设计模块信息
type PredesignedModuleInfo struct {
ModuleName string `json:"moduleName"` // 模块名称
PackageName string `json:"packageName"` // 包名
Template string `json:"template"` // 模板类型
FilePaths []string `json:"filePaths"` // 文件路径列表
Description string `json:"description"` // 描述
}
// CleanupInfo 清理信息
type CleanupInfo struct {
DeletedPackages []string `json:"deletedPackages"` // 已删除的包
DeletedModules []string `json:"deletedModules"` // 已删除的模块
CleanupMessage string `json:"cleanupMessage"` // 清理消息
}
// New 创建KRA分析器工具
func (g *KRAAnalyzer) New() mcp.Tool {
return mcp.NewTool("kra_analyze",
mcp.WithDescription("返回当前系统中有效的包和模块信息,并分析用户需求是否需要创建新的包、模块和字典。同时检查并清理空包,确保系统整洁。"),
mcp.WithString("requirement",
mcp.Description("用户需求描述,用于分析是否需要创建新的包和模块"),
mcp.Required(),
),
)
}
// Handle 处理分析请求
func (g *KRAAnalyzer) Handle(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
if g.db == nil {
return nil, errors.New("数据库连接未初始化")
}
// 解析请求参数
requirementStr, ok := request.GetArguments()["requirement"].(string)
if !ok || requirementStr == "" {
return nil, errors.New("参数错误:requirement 必须是非空字符串")
}
// 创建分析请求
analyzeReq := AnalyzeRequest{
Requirement: requirementStr,
}
// 执行分析逻辑
response, err := g.performAnalysis(ctx, analyzeReq)
if err != nil {
return nil, fmt.Errorf("分析失败: %v", err)
}
// 序列化响应
responseJSON, err := json.Marshal(response)
if err != nil {
return nil, fmt.Errorf("序列化响应失败: %v", err)
}
return &mcp.CallToolResult{
Content: []mcp.Content{
mcp.NewTextContent(string(responseJSON)),
},
}, nil
}
// performAnalysis 执行分析逻辑
func (g *KRAAnalyzer) performAnalysis(ctx context.Context, req AnalyzeRequest) (*AnalyzeResponse, error) {
// 1. 获取数据库中的包信息
var packages []model.SysAutoCodePackage
if err := g.db.Find(&packages).Error; err != nil {
return nil, fmt.Errorf("获取包信息失败: %v", err)
}
// 2. 获取历史记录
var histories []model.SysAutoCodeHistory
if err := g.db.Find(&histories).Error; err != nil {
return nil, fmt.Errorf("获取历史记录失败: %v", err)
}
// 3. 检查空包并进行清理
cleanupInfo := &CleanupInfo{
DeletedPackages: []string{},
DeletedModules: []string{},
}
var validPackages []model.SysAutoCodePackage
var emptyPackageHistoryIDs []uint
for _, pkg := range packages {
isEmpty, err := g.isPackageFolderEmpty(pkg.PackageName, pkg.Template)
if err != nil {
if g.logger != nil {
g.logger.Warn(fmt.Sprintf("检查包 %s 是否为空时出错: %v", pkg.PackageName, err))
}
continue
}
if isEmpty {
// 删除空包文件夹
if err := g.removeEmptyPackageFolder(pkg.PackageName, pkg.Template); err != nil {
if g.logger != nil {
g.logger.Warn(fmt.Sprintf("删除空包文件夹 %s 失败: %v", pkg.PackageName, err))
}
} else {
cleanupInfo.DeletedPackages = append(cleanupInfo.DeletedPackages, pkg.PackageName)
}
// 删除数据库记录
if err := g.db.Delete(&pkg).Error; err != nil {
if g.logger != nil {
g.logger.Warn(fmt.Sprintf("删除包数据库记录 %s 失败: %v", pkg.PackageName, err))
}
}
// 收集相关的历史记录ID
for _, history := range histories {
if history.Package == pkg.PackageName {
emptyPackageHistoryIDs = append(emptyPackageHistoryIDs, uint(history.ID))
cleanupInfo.DeletedModules = append(cleanupInfo.DeletedModules, history.StructName)
}
}
} else {
validPackages = append(validPackages, pkg)
}
}
// 5. 清理空包相关的历史记录和脏历史记录
var dirtyHistoryIDs []uint
for _, history := range histories {
// 检查是否为空包相关的历史记录
for _, emptyID := range emptyPackageHistoryIDs {
if uint(history.ID) == emptyID {
dirtyHistoryIDs = append(dirtyHistoryIDs, uint(history.ID))
break
}
}
}
// 删除脏历史记录
if len(dirtyHistoryIDs) > 0 {
if err := g.db.Delete(&model.SysAutoCodeHistory{}, "id IN ?", dirtyHistoryIDs).Error; err != nil {
if g.logger != nil {
g.logger.Warn(fmt.Sprintf("删除脏历史记录失败: %v", err))
}
} else {
if g.logger != nil {
g.logger.Info(fmt.Sprintf("成功删除 %d 条脏历史记录", len(dirtyHistoryIDs)))
}
}
// 清理相关的API和菜单记录
if err := g.cleanupRelatedApiAndMenus(dirtyHistoryIDs); err != nil {
if g.logger != nil {
g.logger.Warn(fmt.Sprintf("清理相关API和菜单记录失败: %v", err))
}
}
}
// 6. 扫描预设计模块
predesignedModules, err := g.scanPredesignedModules()
if err != nil {
if g.logger != nil {
g.logger.Warn(fmt.Sprintf("扫描预设计模块失败: %v", err))
}
predesignedModules = []PredesignedModuleInfo{} // 设置为空列表,不影响主流程
}
// 7. 过滤掉与已删除包相关的模块
filteredModules := []PredesignedModuleInfo{}
for _, module := range predesignedModules {
isDeleted := false
for _, deletedPkg := range cleanupInfo.DeletedPackages {
if module.PackageName == deletedPkg {
isDeleted = true
break
}
}
if !isDeleted {
filteredModules = append(filteredModules, module)
}
}
// 8. 构建分析结果消息
var analysisMessage strings.Builder
if len(cleanupInfo.DeletedPackages) > 0 || len(cleanupInfo.DeletedModules) > 0 {
analysisMessage.WriteString("**系统清理完成**\n\n")
if len(cleanupInfo.DeletedPackages) > 0 {
analysisMessage.WriteString(fmt.Sprintf("- 删除了 %d 个空包: %s\n", len(cleanupInfo.DeletedPackages), strings.Join(cleanupInfo.DeletedPackages, ", ")))
}
if len(cleanupInfo.DeletedModules) > 0 {
analysisMessage.WriteString(fmt.Sprintf("- 删除了 %d 个相关模块: %s\n", len(cleanupInfo.DeletedModules), strings.Join(cleanupInfo.DeletedModules, ", ")))
}
analysisMessage.WriteString("\n")
cleanupInfo.CleanupMessage = analysisMessage.String()
}
analysisMessage.WriteString(" **分析结果**\n\n")
analysisMessage.WriteString(fmt.Sprintf("- **现有包数量**: %d\n", len(validPackages)))
analysisMessage.WriteString(fmt.Sprintf("- **预设计模块数量**: %d\n\n", len(filteredModules)))
// 9. 转换包信息
existingPackages := make([]PackageInfo, len(validPackages))
for i, pkg := range validPackages {
existingPackages[i] = PackageInfo{
PackageName: pkg.PackageName,
Template: pkg.Template,
Label: pkg.Label,
Desc: pkg.Desc,
Module: pkg.Module,
IsEmpty: false, // 已经过滤掉空包
}
}
dictionaries := []DictionaryPre{} // 这里可以根据需要填充字典信息
err = g.db.Table("sys_dictionaries").Find(&dictionaries, "deleted_at is null").Error
if err != nil {
if g.logger != nil {
g.logger.Warn(fmt.Sprintf("获取字典信息失败: %v", err))
}
dictionaries = []DictionaryPre{} // 设置为空列表,不影响主流程
}
// 10. 构建响应
response := &AnalyzeResponse{
ExistingPackages: existingPackages,
PredesignedModules: filteredModules,
Dictionaries: dictionaries,
}
return response, nil
}
// isPackageFolderEmpty 检查包文件夹是否为空
func (g *KRAAnalyzer) isPackageFolderEmpty(packageName, template string) (bool, error) {
if g.config == nil {
return true, nil // 配置未设置,视为空
}
// 根据模板类型确定基础路径
var basePath string
if template == "plugin" {
basePath = filepath.Join(g.config.Root, g.config.Server, "plugin", packageName)
} else {
basePath = filepath.Join(g.config.Root, g.config.Server, "api", "v1", packageName)
}
// 检查文件夹是否存在
if _, err := os.Stat(basePath); os.IsNotExist(err) {
return true, nil // 文件夹不存在,视为空
} else if err != nil {
return false, err // 其他错误
}
// 递归检查是否有.go文件
return g.hasGoFilesRecursive(basePath)
}
// hasGoFilesRecursive 递归检查目录及其子目录中是否有.go文件
func (g *KRAAnalyzer) hasGoFilesRecursive(dirPath string) (bool, error) {
entries, err := os.ReadDir(dirPath)
if err != nil {
return true, err // 读取失败,返回空
}
// 检查当前目录下的.go文件
for _, entry := range entries {
if !entry.IsDir() && strings.HasSuffix(entry.Name(), ".go") {
return false, nil // 找到.go文件,不为空
}
}
// 递归检查子目录
for _, entry := range entries {
if entry.IsDir() {
subDirPath := filepath.Join(dirPath, entry.Name())
isEmpty, err := g.hasGoFilesRecursive(subDirPath)
if err != nil {
continue // 忽略子目录的错误,继续检查其他目录
}
if !isEmpty {
return false, nil // 子目录中找到.go文件,不为空
}
}
}
return true, nil // 没有找到.go文件,为空
}
// removeEmptyPackageFolder 删除空包文件夹
func (g *KRAAnalyzer) removeEmptyPackageFolder(packageName, template string) error {
if g.config == nil {
return nil // 配置未设置,跳过
}
var basePath string
if template == "plugin" {
basePath = filepath.Join(g.config.Root, g.config.Server, "plugin", packageName)
} else {
// 对于package类型,需要删除多个目录
paths := []string{
filepath.Join(g.config.Root, g.config.Server, "api", "v1", packageName),
filepath.Join(g.config.Root, g.config.Server, "model", packageName),
filepath.Join(g.config.Root, g.config.Server, "router", packageName),
filepath.Join(g.config.Root, g.config.Server, "service", packageName),
}
for _, path := range paths {
if err := g.removeDirectoryIfExists(path); err != nil {
return err
}
}
return nil
}
return g.removeDirectoryIfExists(basePath)
}
// removeDirectoryIfExists 删除目录(如果存在)
func (g *KRAAnalyzer) removeDirectoryIfExists(dirPath string) error {
if _, err := os.Stat(dirPath); os.IsNotExist(err) {
return nil // 目录不存在,无需删除
} else if err != nil {
return err // 其他错误
}
return os.RemoveAll(dirPath)
}
// cleanupRelatedApiAndMenus 清理相关的API和菜单记录
func (g *KRAAnalyzer) cleanupRelatedApiAndMenus(historyIDs []uint) error {
if len(historyIDs) == 0 {
return nil
}
// 这里可以根据需要实现具体的API和菜单清理逻辑
// 由于涉及到具体的业务逻辑,这里只做日志记录
if g.logger != nil {
g.logger.Info(fmt.Sprintf("清理历史记录ID %v 相关的API和菜单记录", historyIDs))
}
return nil
}
// scanPredesignedModules 扫描预设计模块
func (g *KRAAnalyzer) scanPredesignedModules() ([]PredesignedModuleInfo, error) {
if g.config == nil {
return nil, errors.New("autocode配置未设置")
}
// 获取autocode配置路径
autocodeRoot := g.config.Root
if autocodeRoot == "" {
return nil, errors.New("autocode根路径未配置")
}
var modules []PredesignedModuleInfo
// 扫描plugin目录
pluginModules, err := g.scanPluginModules(filepath.Join(autocodeRoot, g.config.Server, "plugin"))
if err != nil {
if g.logger != nil {
g.logger.Warn(fmt.Sprintf("扫描plugin模块失败: %v", err))
}
} else {
modules = append(modules, pluginModules...)
}
// 扫描model目录
modelModules, err := g.scanModelModules(filepath.Join(autocodeRoot, g.config.Server, "model"))
if err != nil {
if g.logger != nil {
g.logger.Warn(fmt.Sprintf("扫描model模块失败: %v", err))
}
} else {
modules = append(modules, modelModules...)
}
return modules, nil
}
// scanPluginModules 扫描插件模块
func (g *KRAAnalyzer) scanPluginModules(pluginDir string) ([]PredesignedModuleInfo, error) {
var modules []PredesignedModuleInfo
if _, err := os.Stat(pluginDir); os.IsNotExist(err) {
return modules, nil // 目录不存在,返回空列表
}
entries, err := os.ReadDir(pluginDir)
if err != nil {
return nil, err
}
for _, entry := range entries {
if entry.IsDir() {
pluginName := entry.Name()
pluginPath := filepath.Join(pluginDir, pluginName)
// 查找model目录
modelDir := filepath.Join(pluginPath, "model")
if _, err := os.Stat(modelDir); err == nil {
// 扫描model目录下的模块
pluginModules, err := g.scanModulesInDirectory(modelDir, pluginName, "plugin")
if err != nil {
if g.logger != nil {
g.logger.Warn(fmt.Sprintf("扫描插件 %s 的模块失败: %v", pluginName, err))
}
continue
}
modules = append(modules, pluginModules...)
}
}
}
return modules, nil
}
// scanModelModules 扫描模型模块
func (g *KRAAnalyzer) scanModelModules(modelDir string) ([]PredesignedModuleInfo, error) {
var modules []PredesignedModuleInfo
if _, err := os.Stat(modelDir); os.IsNotExist(err) {
return modules, nil // 目录不存在,返回空列表
}
entries, err := os.ReadDir(modelDir)
if err != nil {
return nil, err
}
for _, entry := range entries {
if entry.IsDir() {
packageName := entry.Name()
packagePath := filepath.Join(modelDir, packageName)
// 扫描包目录下的模块
packageModules, err := g.scanModulesInDirectory(packagePath, packageName, "package")
if err != nil {
if g.logger != nil {
g.logger.Warn(fmt.Sprintf("扫描包 %s 的模块失败: %v", packageName, err))
}
continue
}
modules = append(modules, packageModules...)
}
}
return modules, nil
}
// scanModulesInDirectory 扫描目录中的模块
func (g *KRAAnalyzer) scanModulesInDirectory(dir, packageName, template string) ([]PredesignedModuleInfo, error) {
var modules []PredesignedModuleInfo
entries, err := os.ReadDir(dir)
if err != nil {
return nil, err
}
for _, entry := range entries {
if !entry.IsDir() && strings.HasSuffix(entry.Name(), ".go") {
moduleName := strings.TrimSuffix(entry.Name(), ".go")
filePath := filepath.Join(dir, entry.Name())
module := PredesignedModuleInfo{
ModuleName: moduleName,
PackageName: packageName,
Template: template,
FilePaths: []string{filePath},
Description: fmt.Sprintf("%s模块中的%s", packageName, moduleName),
}
modules = append(modules, module)
}
}
return modules, nil
}
+755
View File
@@ -0,0 +1,755 @@
package mcpTool
import (
"context"
"encoding/json"
"errors"
"fmt"
"strings"
"github.com/mark3labs/mcp-go/mcp"
"go.uber.org/zap"
"gorm.io/gorm"
"kra/internal/data/model"
)
// 注册工具
func init() {
RegisterTool(&KRAExecutor{})
}
// KRAExecutor KRA代码生成器
type KRAExecutor struct {
db *gorm.DB
logger *zap.Logger
config *AutoCodeConfig
}
// SetDependencies 设置依赖
func (g *KRAExecutor) SetDependencies(db *gorm.DB, logger *zap.Logger, config *AutoCodeConfig) {
g.db = db
g.logger = logger
g.config = config
}
// ExecuteRequest 执行请求结构
type ExecuteRequest struct {
ExecutionPlan ExecutionPlan `json:"executionPlan"` // 执行计划
Requirement string `json:"requirement"` // 原始需求(可选,用于日志记录)
}
// ExecuteResponse 执行响应结构
type ExecuteResponse struct {
Success bool `json:"success"`
Message string `json:"message"`
PackageID uint `json:"packageId,omitempty"`
HistoryID uint `json:"historyId,omitempty"`
Paths map[string]string `json:"paths,omitempty"`
GeneratedPaths []string `json:"generatedPaths,omitempty"`
NextActions []string `json:"nextActions,omitempty"`
}
// ExecutionPlan 执行计划结构
type ExecutionPlan struct {
PackageName string `json:"packageName"`
PackageType string `json:"packageType"` // "plugin" 或 "package"
NeedCreatedPackage bool `json:"needCreatedPackage"`
NeedCreatedModules bool `json:"needCreatedModules"`
NeedCreatedDictionaries bool `json:"needCreatedDictionaries"`
PackageInfo *SysAutoCodePackageCreate `json:"packageInfo,omitempty"`
ModulesInfo []*AutoCode `json:"modulesInfo,omitempty"`
Paths map[string]string `json:"paths,omitempty"`
DictionariesInfo []*DictionaryGenerateRequest `json:"dictionariesInfo,omitempty"`
}
// SysAutoCodePackageCreate 包创建请求
type SysAutoCodePackageCreate struct {
Desc string `json:"desc"`
Label string `json:"label"`
Template string `json:"template"`
PackageName string `json:"packageName"`
}
// AutoCode 自动代码生成配置
type AutoCode struct {
Package string `json:"package"`
TableName string `json:"tableName"`
BusinessDB string `json:"businessDB"`
StructName string `json:"structName"`
PackageName string `json:"packageName"`
Description string `json:"description"`
Abbreviation string `json:"abbreviation"`
HumpPackageName string `json:"humpPackageName"`
GvaModel bool `json:"gvaModel"`
AutoMigrate bool `json:"autoMigrate"`
AutoCreateResource bool `json:"autoCreateResource"`
AutoCreateApiToSql bool `json:"autoCreateApiToSql"`
AutoCreateMenuToSql bool `json:"autoCreateMenuToSql"`
AutoCreateBtnAuth bool `json:"autoCreateBtnAuth"`
OnlyTemplate bool `json:"onlyTemplate"`
IsTree bool `json:"isTree"`
TreeJson string `json:"treeJson"`
IsAdd bool `json:"isAdd"`
GenerateWeb bool `json:"generateWeb"`
GenerateServer bool `json:"generateServer"`
Fields []*AutoField `json:"fields"`
}
// AutoField 自动代码字段配置
type AutoField struct {
FieldName string `json:"fieldName"`
FieldDesc string `json:"fieldDesc"`
FieldType string `json:"fieldType"`
FieldJson string `json:"fieldJson"`
DataTypeLong string `json:"dataTypeLong"`
Comment string `json:"comment"`
ColumnName string `json:"columnName"`
FieldSearchType string `json:"fieldSearchType"`
FieldSearchHide bool `json:"fieldSearchHide"`
DictType string `json:"dictType"`
Form bool `json:"form"`
Table bool `json:"table"`
Desc bool `json:"desc"`
Excel bool `json:"excel"`
Require bool `json:"require"`
DefaultValue string `json:"defaultValue"`
ErrorText string `json:"errorText"`
Clearable bool `json:"clearable"`
Sort bool `json:"sort"`
PrimaryKey bool `json:"primaryKey"`
DataSource *DataSource `json:"dataSource"`
CheckDataSource bool `json:"checkDataSource"`
FieldIndexType string `json:"fieldIndexType"`
}
// DataSource 数据源配置
type DataSource struct {
DBName string `json:"dbName"`
Table string `json:"table"`
Label string `json:"label"`
Value string `json:"value"`
Association int `json:"association"`
HasDeletedAt bool `json:"hasDeletedAt"`
}
// DictOptionItem 字典选项
type DictOptionItem struct {
Label string `json:"label"`
Value string `json:"value"`
Sort int `json:"sort"`
}
// New 创建KRA代码生成执行器工具
func (g *KRAExecutor) New() mcp.Tool {
return mcp.NewTool("kra_execute",
mcp.WithDescription(`**KRA代码生成执行器:直接执行代码生成,无需确认步骤**
**核心功能:**
- 根据需求分析和当前的包信息判断是否调用,如果需要调用,则根据入参描述生成json,用于直接生成代码
- 支持批量创建多个模块
- 自动创建包、模块、字典等
- 移除了确认步骤,提高执行效率
**使用场景:**
- 在kra_analyze获取了当前的包信息和字典信息之后,如果已经包含了可以使用的包和模块,那就不要调用本mcp
- 根据分析结果直接生成代码
- 适用于自动化代码生成流程
**批量创建功能:**
- 支持在单个ExecutionPlan中创建多个模块
- modulesInfo字段为数组,可包含多个模块配置
- 一次性处理多个模块的创建和字典生成
**新功能:自动字典创建**
- 当结构体字段使用了字典类型(dictType不为空)时,系统会自动检查字典是否存在
- 如果字典不存在,会自动创建对应的字典及默认的字典详情项
- 字典创建包括:字典主表记录和默认的选项值(选项1、选项2等)
**重要限制:**
- 当needCreatedModules=true时,模块创建会自动生成API和菜单,因此不应再调用api_creator和menu_creator工具
- 只有在单独创建API或菜单(不涉及模块创建)时才使用api_creator和menu_creator工具`),
mcp.WithObject("executionPlan",
mcp.Description("执行计划,包含包信息、模块与字典信息"),
mcp.Required(),
mcp.Properties(map[string]interface{}{
"packageName": map[string]interface{}{
"type": "string",
"description": "包名(小写开头)",
},
"packageType": map[string]interface{}{
"type": "string",
"description": "package 或 plugin",
"enum": []string{"package", "plugin"},
},
"needCreatedPackage": map[string]interface{}{
"type": "boolean",
"description": "是否需要创建包",
},
"needCreatedModules": map[string]interface{}{
"type": "boolean",
"description": "是否需要创建模块",
},
"needCreatedDictionaries": map[string]interface{}{
"type": "boolean",
"description": "是否需要创建字典",
},
}),
mcp.AdditionalProperties(true),
),
mcp.WithString("requirement",
mcp.Description("原始需求描述(可选,用于日志记录)"),
),
)
}
// Handle 处理执行请求(移除确认步骤)
func (g *KRAExecutor) Handle(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
if g.db == nil {
return nil, errors.New("数据库连接未初始化")
}
executionPlanData, ok := request.GetArguments()["executionPlan"]
if !ok {
return nil, errors.New("参数错误:executionPlan 必须提供")
}
// 解析执行计划
planJSON, err := json.Marshal(executionPlanData)
if err != nil {
return nil, fmt.Errorf("解析执行计划失败: %v", err)
}
var plan ExecutionPlan
err = json.Unmarshal(planJSON, &plan)
if err != nil {
return nil, fmt.Errorf("解析执行计划失败: %v\n\n请确保ExecutionPlan格式正确,参考工具描述中的结构体格式要求", err)
}
// 验证执行计划的完整性
if err := g.validateExecutionPlan(&plan); err != nil {
return nil, fmt.Errorf("执行计划验证失败: %v", err)
}
// 获取原始需求(可选)
var originalRequirement string
if reqData, ok := request.GetArguments()["requirement"]; ok {
if reqStr, ok := reqData.(string); ok {
originalRequirement = reqStr
}
}
// 直接执行创建操作(无确认步骤)
result := g.executeCreation(ctx, &plan)
// 如果执行成功且有原始需求,提供代码复检建议
var reviewMessage string
if result.Success && originalRequirement != "" {
if g.logger != nil {
g.logger.Info("执行完成,返回生成的文件路径供AI进行代码复检...")
}
// 构建文件路径信息供AI使用
var pathsInfo []string
for _, path := range result.GeneratedPaths {
pathsInfo = append(pathsInfo, fmt.Sprintf("- %s", path))
}
reviewMessage = fmt.Sprintf("\n\n📁 已生成以下文件:\n%s\n\n💡 提示:可以检查生成的代码是否满足原始需求。", strings.Join(pathsInfo, "\n"))
} else if originalRequirement == "" {
reviewMessage = "\n\n💡 提示:如需代码复检,请提供原始需求描述。"
}
// 序列化响应
response := ExecuteResponse{
Success: result.Success,
Message: result.Message,
PackageID: result.PackageID,
HistoryID: result.HistoryID,
Paths: result.Paths,
GeneratedPaths: result.GeneratedPaths,
NextActions: result.NextActions,
}
responseJSON, err := json.MarshalIndent(response, "", " ")
if err != nil {
return nil, fmt.Errorf("序列化结果失败: %v", err)
}
return &mcp.CallToolResult{
Content: []mcp.Content{
mcp.NewTextContent(fmt.Sprintf("执行结果:\n\n%s%s", string(responseJSON), reviewMessage)),
},
}, nil
}
// validateExecutionPlan 验证执行计划的完整性
func (g *KRAExecutor) validateExecutionPlan(plan *ExecutionPlan) error {
// 验证基本字段
if plan.PackageName == "" {
return errors.New("packageName 不能为空")
}
if plan.PackageType != "package" && plan.PackageType != "plugin" {
return errors.New("packageType 必须是 'package' 或 'plugin'")
}
// 验证packageType和template字段的一致性
if plan.NeedCreatedPackage && plan.PackageInfo != nil {
if plan.PackageType != plan.PackageInfo.Template {
return errors.New("packageType 和 packageInfo.template 必须保持一致")
}
}
// 验证包信息
if plan.NeedCreatedPackage {
if plan.PackageInfo == nil {
return errors.New("当 needCreatedPackage=true 时,packageInfo 不能为空")
}
if plan.PackageInfo.PackageName == "" {
return errors.New("packageInfo.packageName 不能为空")
}
if plan.PackageInfo.Template != "package" && plan.PackageInfo.Template != "plugin" {
return errors.New("packageInfo.template 必须是 'package' 或 'plugin'")
}
if plan.PackageInfo.Label == "" {
return errors.New("packageInfo.label 不能为空")
}
if plan.PackageInfo.Desc == "" {
return errors.New("packageInfo.desc 不能为空")
}
}
// 验证模块信息(批量验证)
if plan.NeedCreatedModules {
if len(plan.ModulesInfo) == 0 {
return errors.New("当 needCreatedModules=true 时,modulesInfo 不能为空")
}
// 遍历验证每个模块
for moduleIndex, moduleInfo := range plan.ModulesInfo {
if moduleInfo.Package == "" {
return fmt.Errorf("模块 %d 的 package 不能为空", moduleIndex+1)
}
if moduleInfo.StructName == "" {
return fmt.Errorf("模块 %d 的 structName 不能为空", moduleIndex+1)
}
if moduleInfo.TableName == "" {
return fmt.Errorf("模块 %d 的 tableName 不能为空", moduleIndex+1)
}
if moduleInfo.Description == "" {
return fmt.Errorf("模块 %d 的 description 不能为空", moduleIndex+1)
}
if moduleInfo.Abbreviation == "" {
return fmt.Errorf("模块 %d 的 abbreviation 不能为空", moduleIndex+1)
}
if moduleInfo.PackageName == "" {
return fmt.Errorf("模块 %d 的 packageName 不能为空", moduleIndex+1)
}
if moduleInfo.HumpPackageName == "" {
return fmt.Errorf("模块 %d 的 humpPackageName 不能为空", moduleIndex+1)
}
// 验证字段信息
if len(moduleInfo.Fields) == 0 {
return fmt.Errorf("模块 %d 的 fields 不能为空,至少需要一个字段", moduleIndex+1)
}
for i, field := range moduleInfo.Fields {
if field.FieldName == "" {
return fmt.Errorf("模块 %d 字段 %d 的 fieldName 不能为空", moduleIndex+1, i+1)
}
// 确保字段名首字母大写
if len(field.FieldName) > 0 {
firstChar := string(field.FieldName[0])
if firstChar >= "a" && firstChar <= "z" {
moduleInfo.Fields[i].FieldName = strings.ToUpper(firstChar) + field.FieldName[1:]
}
}
if field.FieldDesc == "" {
return fmt.Errorf("模块 %d 字段 %d 的 fieldDesc 不能为空", moduleIndex+1, i+1)
}
if field.FieldType == "" {
return fmt.Errorf("模块 %d 字段 %d 的 fieldType 不能为空", moduleIndex+1, i+1)
}
if field.FieldJson == "" {
return fmt.Errorf("模块 %d 字段 %d 的 fieldJson 不能为空", moduleIndex+1, i+1)
}
if field.ColumnName == "" {
return fmt.Errorf("模块 %d 字段 %d 的 columnName 不能为空", moduleIndex+1, i+1)
}
// 验证字段类型
validFieldTypes := []string{"string", "int", "int64", "float64", "bool", "time.Time", "enum", "picture", "video", "file", "pictures", "array", "richtext", "json"}
validType := false
for _, validFieldType := range validFieldTypes {
if field.FieldType == validFieldType {
validType = true
break
}
}
if !validType {
return fmt.Errorf("模块 %d 字段 %d 的 fieldType '%s' 不支持,支持的类型:%v", moduleIndex+1, i+1, field.FieldType, validFieldTypes)
}
// 验证搜索类型(如果设置了)
if field.FieldSearchType != "" {
validSearchTypes := []string{"=", "!=", ">", ">=", "<", "<=", "LIKE", "BETWEEN", "IN", "NOT IN"}
validSearchType := false
for _, validType := range validSearchTypes {
if field.FieldSearchType == validType {
validSearchType = true
break
}
}
if !validSearchType {
return fmt.Errorf("模块 %d 字段 %d 的 fieldSearchType '%s' 不支持,支持的类型:%v", moduleIndex+1, i+1, field.FieldSearchType, validSearchTypes)
}
}
// 验证 dataSource 字段配置
if field.DataSource != nil {
associationValue := field.DataSource.Association
// 当 association 为 2(一对多关联)时,强制修改 fieldType 为 array
if associationValue == 2 {
if field.FieldType != "array" {
if g.logger != nil {
g.logger.Info(fmt.Sprintf("模块 %d 字段 %d:检测到一对多关联(association=2),自动将 fieldType 从 '%s' 修改为 'array'", moduleIndex+1, i+1, field.FieldType))
}
moduleInfo.Fields[i].FieldType = "array"
}
}
// 验证 association 值的有效性
if associationValue != 1 && associationValue != 2 {
return fmt.Errorf("模块 %d 字段 %d 的 dataSource.association 必须是 1(一对一)或 2(一对多)", moduleIndex+1, i+1)
}
}
}
// 验证主键设置
if !moduleInfo.GvaModel {
// 当不使用默认模型时,必须有且仅有一个字段设置为主键
primaryKeyCount := 0
for _, field := range moduleInfo.Fields {
if field.PrimaryKey {
primaryKeyCount++
}
}
if primaryKeyCount == 0 {
return fmt.Errorf("模块 %d:当 gvaModel=false 时,必须有一个字段的 primaryKey=true", moduleIndex+1)
}
if primaryKeyCount > 1 {
return fmt.Errorf("模块 %d:当 gvaModel=false 时,只能有一个字段的 primaryKey=true", moduleIndex+1)
}
} else {
// 当使用默认模型时,所有字段的primaryKey都应该为false
for i, field := range moduleInfo.Fields {
if field.PrimaryKey {
return fmt.Errorf("模块 %d:当 gvaModel=true 时,字段 %d 的 primaryKey 应该为 false,系统会自动创建ID主键", moduleIndex+1, i+1)
}
}
}
}
}
return nil
}
// executeCreation 执行创建操作
func (g *KRAExecutor) executeCreation(ctx context.Context, plan *ExecutionPlan) *ExecuteResponse {
result := &ExecuteResponse{
Success: false,
Paths: make(map[string]string),
GeneratedPaths: []string{}, // 初始化生成文件路径列表
}
// 无论如何都先构建目录结构信息,确保paths始终返回
result.Paths = g.buildDirectoryStructure(plan)
// 记录预期生成的文件路径
result.GeneratedPaths = g.collectExpectedFilePaths(plan)
if !plan.NeedCreatedModules {
result.Success = true
result.Message += "已列出当前功能所涉及的目录结构信息; 请在paths中查看; 并且在对应指定文件中实现相关的业务逻辑; "
return result
}
// 创建包(如果需要)
if plan.NeedCreatedPackage && plan.PackageInfo != nil {
pkg := model.SysAutoCodePackage{
PackageName: plan.PackageInfo.PackageName,
Template: plan.PackageInfo.Template,
Label: plan.PackageInfo.Label,
Desc: plan.PackageInfo.Desc,
}
err := g.db.Create(&pkg).Error
if err != nil {
result.Message = fmt.Sprintf("创建包失败: %v", err)
// 即使创建包失败,也要返回paths信息
return result
}
result.Message += "包创建成功; "
}
// 创建指定字典(如果需要)
if plan.NeedCreatedDictionaries && len(plan.DictionariesInfo) > 0 {
dictResult := g.createDictionariesFromInfo(ctx, plan.DictionariesInfo)
result.Message += dictResult
}
// 批量创建模块(如果需要)- 注意:这里只是记录,实际代码生成需要调用模板服务
if plan.NeedCreatedModules && len(plan.ModulesInfo) > 0 {
for _, moduleInfo := range plan.ModulesInfo {
// 创建历史记录
history := model.SysAutoCodeHistory{
Package: moduleInfo.Package,
TblName: moduleInfo.TableName,
StructName: moduleInfo.StructName,
Description: moduleInfo.Description,
Abbreviation: moduleInfo.Abbreviation,
}
err := g.db.Create(&history).Error
if err != nil {
result.Message += fmt.Sprintf("创建模块 %s 历史记录失败: %v; ", moduleInfo.StructName, err)
continue
}
result.Message += fmt.Sprintf("模块 %s 记录创建成功; ", moduleInfo.StructName)
}
result.Message += fmt.Sprintf("批量创建完成,共处理 %d 个模块; ", len(plan.ModulesInfo))
// 添加重要提醒:不要使用其他MCP工具
result.Message += "\n\n⚠️ 重要提醒:\n"
result.Message += "模块创建已完成,API和菜单已自动生成。请不要再调用以下MCP工具:\n"
result.Message += "- api_creator:API权限已在模块创建时自动生成\n"
result.Message += "- menu_creator:前端菜单已在模块创建时自动生成\n"
result.Message += "如需修改API或菜单,请直接在系统管理界面中进行配置。\n"
}
result.Message += "已构建目录结构信息; "
result.Success = true
if result.Message == "" {
result.Message = "执行计划完成"
}
return result
}
// buildDirectoryStructure 构建目录结构信息
func (g *KRAExecutor) buildDirectoryStructure(plan *ExecutionPlan) map[string]string {
paths := make(map[string]string)
// 获取配置信息
var rootPath, serverPath, webPath, moduleName string
if g.config != nil {
rootPath = g.config.Root
serverPath = g.config.Server
webPath = g.config.Web
moduleName = g.config.Module
} else {
// 默认配置
rootPath = "."
serverPath = "server"
webPath = "web"
moduleName = "kra"
}
// 如果计划中有包名,使用计划中的包名,否则使用默认
packageName := "example"
if plan.PackageName != "" {
packageName = plan.PackageName
}
// 如果计划中有模块信息,获取第一个模块的结构名作为默认值
structName := "ExampleStruct"
if len(plan.ModulesInfo) > 0 && plan.ModulesInfo[0].StructName != "" {
structName = plan.ModulesInfo[0].StructName
}
// 根据包类型构建不同的路径结构
packageType := plan.PackageType
if packageType == "" {
packageType = "package" // 默认为package模式
}
// 构建服务端路径
if serverPath != "" {
serverBasePath := fmt.Sprintf("%s/%s", rootPath, serverPath)
if packageType == "plugin" {
// Plugin 模式:所有文件都在 /plugin/packageName/ 目录下
plugingBasePath := fmt.Sprintf("%s/plugin/%s", serverBasePath, packageName)
paths["api"] = fmt.Sprintf("%s/api", plugingBasePath)
paths["service"] = fmt.Sprintf("%s/service", plugingBasePath)
paths["model"] = fmt.Sprintf("%s/model", plugingBasePath)
paths["router"] = fmt.Sprintf("%s/router", plugingBasePath)
paths["request"] = fmt.Sprintf("%s/model/request", plugingBasePath)
paths["response"] = fmt.Sprintf("%s/model/response", plugingBasePath)
paths["plugin_main"] = fmt.Sprintf("%s/main.go", plugingBasePath)
paths["plugin_config"] = fmt.Sprintf("%s/plugin.go", plugingBasePath)
paths["plugin_initialize"] = fmt.Sprintf("%s/initialize", plugingBasePath)
} else {
// Package 模式:传统的目录结构
paths["api"] = fmt.Sprintf("%s/api/v1/%s", serverBasePath, packageName)
paths["service"] = fmt.Sprintf("%s/service/%s", serverBasePath, packageName)
paths["model"] = fmt.Sprintf("%s/model/%s", serverBasePath, packageName)
paths["router"] = fmt.Sprintf("%s/router/%s", serverBasePath, packageName)
paths["request"] = fmt.Sprintf("%s/model/%s/request", serverBasePath, packageName)
paths["response"] = fmt.Sprintf("%s/model/%s/response", serverBasePath, packageName)
}
}
// 构建前端路径
if webPath != "" {
webBasePath := fmt.Sprintf("%s/%s", rootPath, webPath)
if packageType == "plugin" {
pluginWebBasePath := fmt.Sprintf("%s/plugin/%s", webBasePath, packageName)
paths["vue_page"] = fmt.Sprintf("%s/view", pluginWebBasePath)
paths["vue_api"] = fmt.Sprintf("%s/api", pluginWebBasePath)
} else {
paths["vue_page"] = fmt.Sprintf("%s/view/%s", webBasePath, packageName)
paths["vue_api"] = fmt.Sprintf("%s/api/%s", webBasePath, packageName)
}
}
// 添加模块信息
paths["module"] = moduleName
paths["package_name"] = packageName
paths["package_type"] = packageType
paths["struct_name"] = structName
paths["root_path"] = rootPath
paths["server_path"] = serverPath
paths["web_path"] = webPath
return paths
}
// collectExpectedFilePaths 收集预期生成的文件路径
func (g *KRAExecutor) collectExpectedFilePaths(plan *ExecutionPlan) []string {
var paths []string
// 获取目录结构
dirPaths := g.buildDirectoryStructure(plan)
// 如果需要创建模块,添加预期的文件路径
if plan.NeedCreatedModules && len(plan.ModulesInfo) > 0 {
for _, moduleInfo := range plan.ModulesInfo {
structName := moduleInfo.StructName
// 后端文件
if apiPath, ok := dirPaths["api"]; ok {
paths = append(paths, fmt.Sprintf("%s/%s.go", apiPath, strings.ToLower(structName)))
}
if servicePath, ok := dirPaths["service"]; ok {
paths = append(paths, fmt.Sprintf("%s/%s.go", servicePath, strings.ToLower(structName)))
}
if modelPath, ok := dirPaths["model"]; ok {
paths = append(paths, fmt.Sprintf("%s/%s.go", modelPath, strings.ToLower(structName)))
}
if routerPath, ok := dirPaths["router"]; ok {
paths = append(paths, fmt.Sprintf("%s/%s.go", routerPath, strings.ToLower(structName)))
}
if requestPath, ok := dirPaths["request"]; ok {
paths = append(paths, fmt.Sprintf("%s/%s.go", requestPath, strings.ToLower(structName)))
}
if responsePath, ok := dirPaths["response"]; ok {
paths = append(paths, fmt.Sprintf("%s/%s.go", responsePath, strings.ToLower(structName)))
}
// 前端文件
if vuePage, ok := dirPaths["vue_page"]; ok {
paths = append(paths, fmt.Sprintf("%s/%s.vue", vuePage, strings.ToLower(structName)))
}
if vueApi, ok := dirPaths["vue_api"]; ok {
paths = append(paths, fmt.Sprintf("%s/%s.js", vueApi, strings.ToLower(structName)))
}
}
}
return paths
}
// checkDictionaryExists 检查字典是否存在
func (g *KRAExecutor) checkDictionaryExists(dictType string) (bool, error) {
var count int64
err := g.db.Model(&model.SysDictionary{}).Where("type = ?", dictType).Count(&count).Error
if err != nil {
return false, err
}
return count > 0, nil
}
// createDictionariesFromInfo 根据 DictionariesInfo 创建字典
func (g *KRAExecutor) createDictionariesFromInfo(ctx context.Context, dictionariesInfo []*DictionaryGenerateRequest) string {
var messages []string
messages = append(messages, fmt.Sprintf("开始创建 %d 个指定字典: ", len(dictionariesInfo)))
for _, dictInfo := range dictionariesInfo {
// 检查字典是否存在
exists, err := g.checkDictionaryExists(dictInfo.DictType)
if err != nil {
messages = append(messages, fmt.Sprintf("检查字典 %s 时出错: %v; ", dictInfo.DictType, err))
continue
}
if !exists {
// 字典不存在,创建字典
dictionary := model.SysDictionary{
Name: dictInfo.DictName,
Type: dictInfo.DictType,
Status: true,
Desc: dictInfo.Description,
}
err = g.db.Create(&dictionary).Error
if err != nil {
messages = append(messages, fmt.Sprintf("创建字典 %s 失败: %v; ", dictInfo.DictType, err))
continue
}
messages = append(messages, fmt.Sprintf("成功创建字典 %s (%s); ", dictInfo.DictType, dictInfo.DictName))
// 获取刚创建的字典ID
var createdDict model.SysDictionary
err = g.db.Where("type = ?", dictInfo.DictType).First(&createdDict).Error
if err != nil {
messages = append(messages, fmt.Sprintf("获取创建的字典失败: %v; ", err))
continue
}
// 创建字典选项
if len(dictInfo.Options) > 0 {
successCount := 0
for _, option := range dictInfo.Options {
dictionaryDetail := model.SysDictionaryDetail{
Label: option.Label,
Value: option.Value,
Status: true,
Sort: int64(option.Sort),
SysDictionaryID: createdDict.ID,
}
err = g.db.Create(&dictionaryDetail).Error
if err != nil {
if g.logger != nil {
g.logger.Warn("创建字典详情项失败", zap.Error(err))
}
} else {
successCount++
}
}
messages = append(messages, fmt.Sprintf("创建了 %d 个字典选项; ", successCount))
}
} else {
messages = append(messages, fmt.Sprintf("字典 %s 已存在,跳过创建; ", dictInfo.DictType))
}
}
return strings.Join(messages, "")
}
+170
View File
@@ -0,0 +1,170 @@
package mcpTool
import (
"context"
"encoding/json"
"errors"
"fmt"
"strings"
"github.com/mark3labs/mcp-go/mcp"
)
// KRAReviewer KRA代码审查工具
type KRAReviewer struct{}
// init 注册工具
func init() {
RegisterTool(&KRAReviewer{})
}
// ReviewRequest 审查请求结构
type ReviewRequest struct {
UserRequirement string `json:"userRequirement"` // 经过requirement_analyze后的用户需求
GeneratedFiles []string `json:"generatedFiles"` // kra_execute创建的文件列表
}
// ReviewResponse 审查响应结构
type ReviewResponse struct {
Success bool `json:"success"` // 是否审查成功
Message string `json:"message"` // 审查结果消息
AdjustmentPrompt string `json:"adjustmentPrompt"` // 调整代码的提示
ReviewDetails string `json:"reviewDetails"` // 详细的审查结果
}
// New 创建KRA代码审查工具
func (g *KRAReviewer) New() mcp.Tool {
return mcp.NewTool("kra_review",
mcp.WithDescription(`**KRA代码审查工具 - 在kra_execute调用后使用**
**核心功能:**
- 接收经过requirement_analyze处理的用户需求和kra_execute生成的文件列表
- 分析生成的代码是否满足用户的原始需求
- 检查是否涉及到关联、交互等复杂功能
- 如果代码不满足需求,提供调整建议和新的prompt
**使用场景:**
- 在kra_execute成功执行后调用
- 用于验证生成的代码是否完整满足用户需求
- 检查模块间的关联关系是否正确实现
- 发现缺失的交互功能或业务逻辑
**工作流程:**
1. 接收用户原始需求和生成的文件列表
2. 分析需求中的关键功能点
3. 检查生成的文件是否覆盖所有功能
4. 识别缺失的关联关系、交互功能等
5. 生成调整建议和新的开发prompt
**输出内容:**
- 审查结果和是否需要调整
- 详细的缺失功能分析
- 针对性的代码调整建议
- 可直接使用的开发prompt
**重要提示:**
- 本工具专门用于代码质量审查,不执行实际的代码修改
- 重点关注模块间关联、用户交互、业务流程完整性
- 提供的调整建议应该具体可执行`),
mcp.WithString("userRequirement",
mcp.Description("经过requirement_analyze处理后的用户需求描述,包含详细的功能要求和字段信息"),
mcp.Required(),
),
mcp.WithString("generatedFiles",
mcp.Description("kra_execute创建的文件列表,JSON字符串格式,包含所有生成的后端和前端文件路径"),
mcp.Required(),
),
)
}
// Handle 处理审查请求
func (g *KRAReviewer) Handle(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
// 获取用户需求
userRequirementData, ok := request.GetArguments()["userRequirement"]
if !ok {
return nil, errors.New("参数错误:userRequirement 必须提供")
}
userRequirement, ok := userRequirementData.(string)
if !ok {
return nil, errors.New("参数错误:userRequirement 必须是字符串类型")
}
// 获取生成的文件列表
generatedFilesData, ok := request.GetArguments()["generatedFiles"]
if !ok {
return nil, errors.New("参数错误:generatedFiles 必须提供")
}
generatedFilesStr, ok := generatedFilesData.(string)
if !ok {
return nil, errors.New("参数错误:generatedFiles 必须是JSON字符串")
}
// 解析JSON字符串为字符串数组
var generatedFiles []string
err := json.Unmarshal([]byte(generatedFilesStr), &generatedFiles)
if err != nil {
return nil, fmt.Errorf("解析generatedFiles失败: %v", err)
}
if len(generatedFiles) == 0 {
return nil, errors.New("参数错误:generatedFiles 不能为空")
}
// 直接生成调整提示,不进行复杂分析
adjustmentPrompt := g.generateAdjustmentPrompt(userRequirement, generatedFiles)
// 构建简化的审查详情
reviewDetails := fmt.Sprintf("📋 **代码审查报告**\n\n **用户原始需求:**\n%s\n\n **已生成文件数量:** %d\n\n **建议进行代码优化和完善**", userRequirement, len(generatedFiles))
// 构建审查结果
reviewResult := &ReviewResponse{
Success: true,
Message: "代码审查完成",
AdjustmentPrompt: adjustmentPrompt,
ReviewDetails: reviewDetails,
}
// 序列化响应
responseJSON, err := json.MarshalIndent(reviewResult, "", " ")
if err != nil {
return nil, fmt.Errorf("序列化审查结果失败: %v", err)
}
return &mcp.CallToolResult{
Content: []mcp.Content{
mcp.NewTextContent(fmt.Sprintf("代码审查结果:\n\n%s", string(responseJSON))),
},
}, nil
}
// generateAdjustmentPrompt 生成调整代码的提示
func (g *KRAReviewer) generateAdjustmentPrompt(userRequirement string, generatedFiles []string) string {
var prompt strings.Builder
prompt.WriteString("🔧 **代码调整指导 Prompt:**\n\n")
prompt.WriteString(fmt.Sprintf("**用户的原始需求为:** %s\n\n", userRequirement))
prompt.WriteString("**经过KRA生成后的文件有如下内容:**\n")
for _, file := range generatedFiles {
prompt.WriteString(fmt.Sprintf("- %s\n", file))
}
prompt.WriteString("\n")
prompt.WriteString("**请帮我优化和完善代码,确保:**\n")
prompt.WriteString("1. 代码完全满足用户的原始需求\n")
prompt.WriteString("2. 完善模块间的关联关系,确保数据一致性\n")
prompt.WriteString("3. 实现所有必要的用户交互功能\n")
prompt.WriteString("4. 保持代码的完整性和可维护性\n")
prompt.WriteString("5. 遵循KRA框架的开发规范和最佳实践\n")
prompt.WriteString("6. 确保前后端功能完整对接\n")
prompt.WriteString("7. 添加必要的错误处理和数据验证\n\n")
prompt.WriteString("8. 如果需要vue路由跳转,请使用 menu_lister获取完整路由表,并且路由跳转使用 router.push({\"name\":从menu_lister中获取的name})\n\n")
prompt.WriteString("9. 如果当前所有的vue页面内容无法满足需求,则自行书写vue文件,并且调用 menu_creator创建菜单记录\n\n")
prompt.WriteString("10. 如果需要API调用,请使用 api_lister获取api表,根据需求调用对应接口\n\n")
prompt.WriteString("11. 如果当前所有API无法满足则自行书写接口,补全前后端代码,并使用 api_creator创建api记录\n\n")
prompt.WriteString("12. 无论前后端都不要随意删除import的内容\n\n")
prompt.WriteString("**请基于用户需求和现有文件,提供完整的代码优化方案。**")
return prompt.String()
}
+285
View File
@@ -0,0 +1,285 @@
package mcpTool
import (
"context"
"encoding/json"
"errors"
"fmt"
"github.com/mark3labs/mcp-go/mcp"
"go.uber.org/zap"
"gorm.io/gorm"
"kra/internal/data/model"
)
// 注册工具
func init() {
RegisterTool(&MenuCreator{})
}
// MenuCreateRequest 菜单创建请求结构
type MenuCreateRequest struct {
ParentId uint `json:"parentId"` // 父菜单ID,0表示根菜单
Path string `json:"path"` // 路由path
Name string `json:"name"` // 路由name
Hidden bool `json:"hidden"` // 是否在列表隐藏
Component string `json:"component"` // 对应前端文件路径
Sort int `json:"sort"` // 排序标记
Title string `json:"title"` // 菜单名
Icon string `json:"icon"` // 菜单图标
KeepAlive bool `json:"keepAlive"` // 是否缓存
DefaultMenu bool `json:"defaultMenu"` // 是否是基础路由
CloseTab bool `json:"closeTab"` // 自动关闭tab
ActiveName string `json:"activeName"` // 高亮菜单
Parameters []MenuParameterRequest `json:"parameters"` // 路由参数
MenuBtn []MenuButtonRequest `json:"menuBtn"` // 菜单按钮
}
// MenuParameterRequest 菜单参数请求结构
type MenuParameterRequest struct {
Type string `json:"type"` // 参数类型:params或query
Key string `json:"key"` // 参数key
Value string `json:"value"` // 参数值
}
// MenuButtonRequest 菜单按钮请求结构
type MenuButtonRequest struct {
Name string `json:"name"` // 按钮名称
Desc string `json:"desc"` // 按钮描述
}
// MenuCreateResponse 菜单创建响应结构
type MenuCreateResponse struct {
Success bool `json:"success"`
Message string `json:"message"`
MenuID uint `json:"menuId"`
Name string `json:"name"`
Path string `json:"path"`
}
// MenuCreator 菜单创建工具
type MenuCreator struct {
db *gorm.DB
logger *zap.Logger
}
// SetDependencies 设置依赖
func (m *MenuCreator) SetDependencies(db *gorm.DB, logger *zap.Logger) {
m.db = db
m.logger = logger
}
// New 创建菜单创建工具
func (m *MenuCreator) New() mcp.Tool {
return mcp.NewTool("create_menu",
mcp.WithDescription(`创建前端菜单记录,用于AI编辑器自动添加前端页面时自动创建对应的菜单项。
**重要限制:**
- 当使用kra_auto_generate工具且needCreatedModules=true时,模块创建会自动生成菜单项,不应调用此工具
- 仅在以下情况使用:1) 单独创建菜单(不涉及模块创建);2) AI编辑器自动添加前端页面时`),
mcp.WithNumber("parentId",
mcp.Description("父菜单ID,0表示根菜单"),
mcp.DefaultNumber(0),
),
mcp.WithString("path",
mcp.Required(),
mcp.Description("路由path,如:userList"),
),
mcp.WithString("name",
mcp.Required(),
mcp.Description("路由name,用于Vue Router,如:userList"),
),
mcp.WithBoolean("hidden",
mcp.Description("是否在菜单列表中隐藏"),
),
mcp.WithString("component",
mcp.Required(),
mcp.Description("对应的前端Vue组件路径,如:view/user/list.vue"),
),
mcp.WithNumber("sort",
mcp.Description("菜单排序号,数字越小越靠前"),
mcp.DefaultNumber(1),
),
mcp.WithString("title",
mcp.Required(),
mcp.Description("菜单显示标题"),
),
mcp.WithString("icon",
mcp.Description("菜单图标名称"),
mcp.DefaultString("menu"),
),
mcp.WithBoolean("keepAlive",
mcp.Description("是否缓存页面"),
),
mcp.WithBoolean("defaultMenu",
mcp.Description("是否是基础路由"),
),
mcp.WithBoolean("closeTab",
mcp.Description("是否自动关闭tab"),
),
mcp.WithString("activeName",
mcp.Description("高亮菜单名称"),
),
mcp.WithString("parameters",
mcp.Description("路由参数JSON字符串,格式:[{\"type\":\"params\",\"key\":\"id\",\"value\":\"1\"}]"),
),
mcp.WithString("menuBtn",
mcp.Description("菜单按钮JSON字符串,格式:[{\"name\":\"add\",\"desc\":\"新增\"}]"),
),
)
}
// Handle 处理菜单创建请求
func (m *MenuCreator) Handle(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
if m.db == nil {
return nil, errors.New("数据库连接未初始化")
}
// 解析请求参数
args := request.GetArguments()
// 必需参数
path, ok := args["path"].(string)
if !ok || path == "" {
return nil, errors.New("path 参数是必需的")
}
name, ok := args["name"].(string)
if !ok || name == "" {
return nil, errors.New("name 参数是必需的")
}
component, ok := args["component"].(string)
if !ok || component == "" {
return nil, errors.New("component 参数是必需的")
}
title, ok := args["title"].(string)
if !ok || title == "" {
return nil, errors.New("title 参数是必需的")
}
// 可选参数
parentId := int64(0)
if val, ok := args["parentId"].(float64); ok {
parentId = int64(val)
}
hidden := false
if val, ok := args["hidden"].(bool); ok {
hidden = val
}
sort := 1
if val, ok := args["sort"].(float64); ok {
sort = int(val)
}
icon := "menu"
if val, ok := args["icon"].(string); ok && val != "" {
icon = val
}
keepAlive := false
if val, ok := args["keepAlive"].(bool); ok {
keepAlive = val
}
defaultMenu := false
if val, ok := args["defaultMenu"].(bool); ok {
defaultMenu = val
}
closeTab := false
if val, ok := args["closeTab"].(bool); ok {
closeTab = val
}
activeName := ""
if val, ok := args["activeName"].(string); ok {
activeName = val
}
// 解析参数和按钮
var parameters []model.SysBaseMenuParameter
if parametersStr, ok := args["parameters"].(string); ok && parametersStr != "" {
var paramReqs []MenuParameterRequest
if err := json.Unmarshal([]byte(parametersStr), &paramReqs); err != nil {
return nil, fmt.Errorf("parameters 参数格式错误: %v", err)
}
for _, param := range paramReqs {
parameters = append(parameters, model.SysBaseMenuParameter{
Type: param.Type,
Key: param.Key,
Value: param.Value,
})
}
}
var menuBtn []model.SysBaseMenuBtn
if menuBtnStr, ok := args["menuBtn"].(string); ok && menuBtnStr != "" {
var btnReqs []MenuButtonRequest
if err := json.Unmarshal([]byte(menuBtnStr), &btnReqs); err != nil {
return nil, fmt.Errorf("menuBtn 参数格式错误: %v", err)
}
for _, btn := range btnReqs {
menuBtn = append(menuBtn, model.SysBaseMenuBtn{
Name: btn.Name,
Desc: btn.Desc,
})
}
}
// 构建菜单对象
menu := model.SysBaseMenu{
ParentID: parentId,
Path: path,
Name: name,
Hidden: hidden,
Component: component,
Sort: int64(sort),
Title: title,
Icon: icon,
KeepAlive: keepAlive,
DefaultMenu: defaultMenu,
CloseTab: closeTab,
ActiveName: activeName,
}
// 创建菜单
err := m.db.Create(&menu).Error
if err != nil {
return nil, fmt.Errorf("创建菜单失败: %v", err)
}
// 获取创建的菜单ID
var createdMenu model.SysBaseMenu
err = m.db.Where("name = ? AND path = ?", name, path).First(&createdMenu).Error
if err != nil && m.logger != nil {
m.logger.Warn("获取创建的菜单ID失败", zap.Error(err))
}
// 构建响应
response := &MenuCreateResponse{
Success: true,
Message: fmt.Sprintf("成功创建菜单 %s", title),
MenuID: uint(createdMenu.ID),
Name: name,
Path: path,
}
resultJSON, err := json.MarshalIndent(response, "", " ")
if err != nil {
return nil, fmt.Errorf("序列化结果失败: %v", err)
}
return &mcp.CallToolResult{
Content: []mcp.Content{
mcp.TextContent{
Type: "text",
Text: fmt.Sprintf("菜单创建结果:\n\n%s", string(resultJSON)),
},
},
}, nil
}
+139
View File
@@ -0,0 +1,139 @@
package mcpTool
import (
"context"
"encoding/json"
"fmt"
"github.com/mark3labs/mcp-go/mcp"
"go.uber.org/zap"
"gorm.io/gorm"
"kra/internal/data/model"
)
// 注册工具
func init() {
RegisterTool(&MenuLister{})
}
// MenuListResponse 菜单列表响应结构
type MenuListResponse struct {
Success bool `json:"success"`
Message string `json:"message"`
Menus []model.SysBaseMenu `json:"menus"`
TotalCount int `json:"totalCount"`
Description string `json:"description"`
}
// MenuLister 菜单列表工具
type MenuLister struct {
db *gorm.DB
logger *zap.Logger
}
// SetDependencies 设置依赖
func (m *MenuLister) SetDependencies(db *gorm.DB, logger *zap.Logger) {
m.db = db
m.logger = logger
}
// New 创建菜单列表工具
func (m *MenuLister) New() mcp.Tool {
return mcp.NewTool("list_all_menus",
mcp.WithDescription(`获取系统中所有菜单信息,包括菜单树结构、路由信息、组件路径等,用于前端编写vue-router时正确跳转
**功能说明:**
- 返回完整的菜单树形结构
- 包含路由配置信息(path、name、component)
- 包含菜单元数据(title、icon、keepAlive等)
- 包含菜单参数和按钮配置
- 支持父子菜单关系展示
**使用场景:**
- 前端路由配置:获取所有菜单信息用于配置vue-router
- 菜单权限管理:了解系统中所有可用的菜单项
- 导航组件开发:构建动态导航菜单
- 系统架构分析:了解系统的菜单结构和页面组织`),
mcp.WithString("_placeholder",
mcp.Description("占位符,防止json schema校验失败"),
),
)
}
// Handle 处理菜单列表请求
func (m *MenuLister) Handle(_ context.Context, _ mcp.CallToolRequest) (*mcp.CallToolResult, error) {
if m.db == nil {
return &mcp.CallToolResult{
Content: []mcp.Content{
mcp.TextContent{
Type: "text",
Text: "数据库连接未初始化",
},
},
IsError: true,
}, nil
}
// 获取所有基础菜单
allMenus, err := m.getAllMenus()
if err != nil {
if m.logger != nil {
m.logger.Error("获取菜单列表失败", zap.Error(err))
}
return &mcp.CallToolResult{
Content: []mcp.Content{
mcp.TextContent{
Type: "text",
Text: fmt.Sprintf("获取菜单列表失败: %v", err),
},
},
IsError: true,
}, nil
}
// 构建返回结果
response := MenuListResponse{
Success: true,
Message: "获取菜单列表成功",
Menus: allMenus,
TotalCount: len(allMenus),
Description: "系统中所有菜单信息的标准列表,包含路由配置和组件信息",
}
// 序列化响应
responseJSON, err := json.MarshalIndent(response, "", " ")
if err != nil {
if m.logger != nil {
m.logger.Error("序列化菜单响应失败", zap.Error(err))
}
return &mcp.CallToolResult{
Content: []mcp.Content{
mcp.TextContent{
Type: "text",
Text: fmt.Sprintf("序列化响应失败: %v", err),
},
},
IsError: true,
}, nil
}
return &mcp.CallToolResult{
Content: []mcp.Content{
mcp.TextContent{
Type: "text",
Text: string(responseJSON),
},
},
}, nil
}
// getAllMenus 获取所有基础菜单
func (m *MenuLister) getAllMenus() ([]model.SysBaseMenu, error) {
var menus []model.SysBaseMenu
err := m.db.Order("sort").Preload("Parameters").Preload("MenuBtn").Find(&menus).Error
if err != nil {
return nil, err
}
return menus, nil
}
+199
View File
@@ -0,0 +1,199 @@
package mcpTool
import (
"context"
"encoding/json"
"errors"
"fmt"
"github.com/mark3labs/mcp-go/mcp"
)
func init() {
RegisterTool(&RequirementAnalyzer{})
}
type RequirementAnalyzer struct{}
// RequirementAnalysisRequest 需求分析请求
type RequirementAnalysisRequest struct {
UserRequirement string `json:"userRequirement"`
}
// RequirementAnalysisResponse 需求分析响应
type RequirementAnalysisResponse struct {
AIPrompt string `json:"aiPrompt"` // 给AI的提示词
}
// New 返回工具注册信息
func (t *RequirementAnalyzer) New() mcp.Tool {
return mcp.NewTool("requirement_analyzer",
mcp.WithDescription(`** 智能需求分析与模块设计工具 - 首选入口工具(最高优先级)**
** 重要提示:这是所有MCP工具的首选入口,请优先使用!**
** 核心能力:**
作为资深系统架构师,智能分析用户需求并自动设计完整的模块架构
** 核心功能:**
1. **智能需求解构**:深度分析用户需求,识别核心业务实体、业务流程、数据关系
2. **自动模块设计**:基于需求分析,智能确定需要多少个模块及各模块功能
3. **字段智能推导**:为每个模块自动设计详细字段,包含数据类型、关联关系、字典需求
4. **架构优化建议**:提供模块拆分、关联设计、扩展性等专业建议
** 输出内容:**
- 模块数量和架构设计
- 每个模块的详细字段清单
- 数据类型和关联关系设计
- 字典需求和类型定义
- 模块间关系图和扩展建议
** 适用场景:**
- 用户需求描述不完整,需要智能补全
- 复杂业务系统的模块架构设计
- 需要专业的数据库设计建议
- 想要快速搭建生产级业务系统
** 推荐工作流:**
requirement_analyzer → kra_analyze → kra_execute → 其他辅助工具
`),
mcp.WithString("userRequirement",
mcp.Required(),
mcp.Description("用户的需求描述,支持自然语言,如:'我要做一个猫舍管理系统,用来录入猫的信息,并且记录每只猫每天的活动信息'"),
),
)
}
// Handle 处理工具调用
func (t *RequirementAnalyzer) Handle(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
userRequirement, ok := request.GetArguments()["userRequirement"].(string)
if !ok || userRequirement == "" {
return nil, errors.New("参数错误:userRequirement 必须是非空字符串")
}
// 分析用户需求
analysisResponse, err := t.analyzeRequirement(userRequirement)
if err != nil {
return nil, fmt.Errorf("需求分析失败: %v", err)
}
// 序列化响应
responseData, err := json.Marshal(analysisResponse)
if err != nil {
return nil, fmt.Errorf("序列化响应失败: %v", err)
}
return &mcp.CallToolResult{
Content: []mcp.Content{
mcp.NewTextContent(string(responseData)),
},
}, nil
}
// analyzeRequirement 分析用户需求 - 专注于AI需求传递
func (t *RequirementAnalyzer) analyzeRequirement(userRequirement string) (*RequirementAnalysisResponse, error) {
// 生成AI提示词 - 这是唯一功能
aiPrompt := t.generateAIPrompt(userRequirement)
return &RequirementAnalysisResponse{
AIPrompt: aiPrompt,
}, nil
}
// generateAIPrompt 生成AI提示词 - 智能分析需求并确定模块结构
func (t *RequirementAnalyzer) generateAIPrompt(userRequirement string) string {
prompt := fmt.Sprintf(`# 智能需求分析与模块设计任务
## 用户原始需求
%s
## 核心任务
你需要作为一个资深的系统架构师,深度分析用户需求,智能设计出完整的模块架构。
## 分析步骤
### 第一步:需求解构分析
请仔细分析用户需求,识别出:
1. **核心业务实体**(如:用户、商品、订单、疫苗、宠物等)
2. **业务流程**(如:注册、购买、记录、管理等)
3. **数据关系**(实体间的关联关系)
4. **功能模块**(需要哪些独立的管理模块)
### 第二步:模块架构设计
基于需求分析,设计出模块架构,格式如下:
**模块1:[模块名称]**
- 功能描述:[该模块的核心功能]
- 主要字段:[列出关键字段,注明数据类型]
- 关联关系:[与其他模块的关系,明确一对一/一对多]
- 字典需求:[需要哪些字典类型]
**模块2:[模块名称]**
- 功能描述:[该模块的核心功能]
- 主要字段:[列出关键字段,注明数据类型]
- 关联关系:[与其他模块的关系]
- 字典需求:[需要哪些字典类型]
**...**
### 第三步:字段详细设计
为每个模块详细设计字段:
#### 模块1字段清单:
- 字段名1 (数据类型) - 字段描述 [是否必填] [关联信息/字典类型]
- 字段名2 (数据类型) - 字段描述 [是否必填] [关联信息/字典类型]
- ...
#### 模块2字段清单:
- 字段名1 (数据类型) - 字段描述 [是否必填] [关联信息/字典类型]
- ...
## 智能分析指导原则
### 模块拆分原则
1. **单一职责**:每个模块只负责一个核心业务实体
2. **数据完整性**:相关数据应该在同一模块中
3. **业务独立性**:模块应该能够独立完成特定业务功能
4. **扩展性考虑**:为未来功能扩展预留空间
### 字段设计原则
1. **必要性**:只包含业务必需的字段
2. **规范性**:遵循数据库设计规范
3. **关联性**:正确识别实体间关系
4. **字典化**:状态、类型等枚举值使用字典
### 关联关系识别
- **一对一**:一个实体只能关联另一个实体的一个记录
- **一对多**:一个实体可以关联另一个实体的多个记录
- **多对多**:通过中间表实现复杂关联
## 特殊场景处理
### 复杂实体识别
当用户提到某个概念时,要判断它是否需要独立模块:
- **字典处理**:简单的常见的状态、类型(如:开关、性别、完成状态等)
- **独立模块**:复杂实体(如:疫苗管理、宠物档案、注射记录)
## 输出要求
### 必须包含的信息
1. **模块数量**:明确需要几个模块
2. **模块关系图**:用文字描述模块间关系
3. **核心字段**:每个模块的关键字段(至少5-10个)
4. **数据类型**:string、int、bool、time.Time、float64等
5. **关联设计**:明确哪些字段是关联字段
6. **字典需求**:列出需要创建的字典类型
### 严格遵循用户输入
- 如果用户提供了具体字段,**必须使用**用户提供的字段
- 如果用户提供了SQL文件,**严格按照**SQL结构设计
- **不要**随意发散,**不要**添加用户未提及的功能
---
**现在请开始深度分析用户需求:"%s"**
请按照上述框架进行系统性分析,确保输出的模块设计既满足当前需求,又具备良好的扩展性。`, userRequirement, userRequirement)
return prompt
}
+1 -1
View File
@@ -6,7 +6,7 @@ import (
"github.com/gin-gonic/gin"
)
// Response 统一响应结构(与 gva 保持一致)
// Response 统一响应结构(与 kra 保持一致)
type Response struct {
Code int `json:"code"`
Data interface{} `json:"data"`
+72
View File
@@ -0,0 +1,72 @@
package timer
import (
"fmt"
"testing"
"time"
"github.com/stretchr/testify/assert"
)
var job = mockJob{}
type mockJob struct{}
func (job mockJob) Run() {
mockFunc()
}
func mockFunc() {
time.Sleep(time.Second)
fmt.Println("1s...")
}
func TestNewTimerTask(t *testing.T) {
tm := NewTimerTask()
_tm := tm.(*timer)
{
_, err := tm.AddTaskByFunc("func", "@every 1s", mockFunc, "测试mockfunc")
assert.Nil(t, err)
_, ok := _tm.cronList["func"]
if !ok {
t.Error("no find func")
}
}
{
_, err := tm.AddTaskByJob("job", "@every 1s", job, "测试job mockfunc")
assert.Nil(t, err)
_, ok := _tm.cronList["job"]
if !ok {
t.Error("no find job")
}
}
{
_, ok := tm.FindCron("func")
if !ok {
t.Error("no find func")
}
_, ok = tm.FindCron("job")
if !ok {
t.Error("no find job")
}
_, ok = tm.FindCron("none")
if ok {
t.Error("find none")
}
}
{
tm.Clear("func")
_, ok := tm.FindCron("func")
if ok {
t.Error("find func")
}
}
{
a := tm.FindCronList()
b, c := tm.FindCron("job")
fmt.Println(a, b, c)
}
}
+63 -29
View File
@@ -1,12 +1,20 @@
package upload
import (
"errors"
"io"
"mime/multipart"
"os"
"path"
"path/filepath"
"strings"
"sync"
"time"
"kra/pkg/utils"
)
var mu sync.Mutex
// Local 本地存储
type Local struct {
StorePath string
@@ -23,44 +31,70 @@ func NewLocal(storePath, basePath string) *Local {
// UploadFile 上传文件
func (l *Local) UploadFile(file *multipart.FileHeader) (string, string, error) {
// 生成文件名和路径
fileName := GenerateFileName(file.Filename)
filePath := GeneratePath(l.BasePath)
fullPath := path.Join(l.StorePath, filePath)
// 读取文件后缀
ext := filepath.Ext(file.Filename)
// 读取文件名并加密
name := strings.TrimSuffix(file.Filename, ext)
name = utils.MD5V([]byte(name))
// 拼接新文件名
filename := name + "_" + time.Now().Format("20060102150405") + ext
// 创建目录
if err := os.MkdirAll(fullPath, os.ModePerm); err != nil {
return "", "", err
// 尝试创建此路径
mkdirErr := os.MkdirAll(l.StorePath, os.ModePerm)
if mkdirErr != nil {
return "", "", errors.New("function os.MkdirAll() failed, err:" + mkdirErr.Error())
}
// 打开源文件
src, err := file.Open()
if err != nil {
return "", "", err
}
defer src.Close()
// 拼接路径和文件名
p := l.StorePath + "/" + filename
filePath := l.BasePath + "/" + filename
// 创建目标文件
dstPath := path.Join(fullPath, fileName)
dst, err := os.Create(dstPath)
if err != nil {
return "", "", err
f, openError := file.Open() // 读取文件
if openError != nil {
return "", "", errors.New("function file.Open() failed, err:" + openError.Error())
}
defer dst.Close()
defer f.Close()
// 复制文件
if _, err = io.Copy(dst, src); err != nil {
return "", "", err
out, createErr := os.Create(p)
if createErr != nil {
return "", "", errors.New("function os.Create() failed, err:" + createErr.Error())
}
defer out.Close()
// 返回访问路径和存储key
accessPath := path.Join("/", filePath, fileName)
key := path.Join(filePath, fileName)
return accessPath, key, nil
_, copyErr := io.Copy(out, f) // 传输(拷贝)文件
if copyErr != nil {
return "", "", errors.New("function io.Copy() failed, err:" + copyErr.Error())
}
return filePath, filename, nil
}
// DeleteFile 删除文件
func (l *Local) DeleteFile(key string) error {
fullPath := path.Join(l.StorePath, key)
return os.Remove(fullPath)
// 检查 key 是否为空
if key == "" {
return errors.New("key不能为空")
}
// 验证 key 是否包含非法字符或尝试访问存储路径之外的文件
if strings.Contains(key, "..") || strings.ContainsAny(key, `\/:*?"<>|`) {
return errors.New("非法的key")
}
p := filepath.Join(l.StorePath, key)
// 检查文件是否存在
if _, err := os.Stat(p); os.IsNotExist(err) {
return errors.New("文件不存在")
}
// 使用文件锁防止并发删除
mu.Lock()
defer mu.Unlock()
err := os.Remove(p)
if err != nil {
return errors.New("文件删除失败: " + err.Error())
}
return nil
}
+82 -47
View File
@@ -1,84 +1,119 @@
package upload
import (
"bytes"
"context"
"errors"
"io"
"mime"
"mime/multipart"
"path"
"path/filepath"
"strings"
"time"
"kra/pkg/utils"
"github.com/minio/minio-go/v7"
"github.com/minio/minio-go/v7/pkg/credentials"
)
var MinioClient *Minio // 优化性能,但是不支持动态配置
// Minio MinIO存储
type Minio struct {
ID string
Secret string
Bucket string
Endpoint string
BasePath string
UseSSL bool
Client *minio.Client
Bucket string
BasePath string
BucketURL string
}
// NewMinio 创建MinIO
func NewMinio(id, secret, bucket, endpoint, basePath string, useSSL bool) *Minio {
func NewMinio(accessKeyID, accessKeySecret, bucket, endpoint, basePath string, useSSL bool) *Minio {
if MinioClient != nil {
return MinioClient
}
return &Minio{
ID: id,
Secret: secret,
Bucket: bucket,
Endpoint: endpoint,
BasePath: basePath,
UseSSL: useSSL,
}
}
func (m *Minio) getClient() (*minio.Client, error) {
return minio.New(m.Endpoint, &minio.Options{
Creds: credentials.NewStaticV4(m.ID, m.Secret, ""),
Secure: m.UseSSL,
// GetMinio 获取MinIO客户端实例
func GetMinio(endpoint, accessKeyID, accessKeySecret, bucketName string, useSSL bool) (*Minio, error) {
if MinioClient != nil {
return MinioClient, nil
}
// Initialize minio client object.
minioClient, err := minio.New(endpoint, &minio.Options{
Creds: credentials.NewStaticV4(accessKeyID, accessKeySecret, ""),
Secure: useSSL,
})
if err != nil {
return nil, err
}
// 尝试创建bucket
err = minioClient.MakeBucket(context.Background(), bucketName, minio.MakeBucketOptions{})
if err != nil {
// Check to see if we already own this bucket (which happens if you run this twice)
exists, errBucketExists := minioClient.BucketExists(context.Background(), bucketName)
if errBucketExists == nil && exists {
// log.Printf("We already own %s\n", bucketName)
} else {
return nil, err
}
}
MinioClient = &Minio{Client: minioClient, Bucket: bucketName}
return MinioClient, nil
}
// UploadFile 上传文件
func (m *Minio) UploadFile(file *multipart.FileHeader) (string, string, error) {
client, err := m.getClient()
f, openError := file.Open()
if openError != nil {
return "", "", errors.New("function file.Open() Failed, err:" + openError.Error())
}
filecontent := bytes.Buffer{}
_, err := io.Copy(&filecontent, f)
if err != nil {
return "", "", err
return "", "", errors.New("读取文件失败, err:" + err.Error())
}
f.Close()
// 对文件名进行加密存储
ext := filepath.Ext(file.Filename)
filename := utils.MD5V([]byte(strings.TrimSuffix(file.Filename, ext))) + ext
var filePathRes string
if m.BasePath == "" {
filePathRes = "uploads" + "/" + time.Now().Format("2006-01-02") + "/" + filename
} else {
filePathRes = m.BasePath + "/" + time.Now().Format("2006-01-02") + "/" + filename
}
// 生成文件名和路径
fileName := GenerateFileName(file.Filename)
filePath := GeneratePath(m.BasePath)
key := path.Join(filePath, fileName)
// 根据文件扩展名检测 MIME 类型
contentType := mime.TypeByExtension(ext)
if contentType == "" {
contentType = "application/octet-stream"
}
// 打开文件
src, err := file.Open()
// 设置超时10分钟
ctx, cancel := context.WithTimeout(context.Background(), time.Minute*10)
defer cancel()
// Upload the file with PutObject 大文件自动切换为分片上传
info, err := m.Client.PutObject(ctx, m.Bucket, filePathRes, &filecontent, file.Size, minio.PutObjectOptions{ContentType: contentType})
if err != nil {
return "", "", err
return "", "", errors.New("上传文件到minio失败, err:" + err.Error())
}
defer src.Close()
// 上传
_, err = client.PutObject(context.Background(), m.Bucket, key, src, file.Size, minio.PutObjectOptions{
ContentType: file.Header.Get("Content-Type"),
})
if err != nil {
return "", "", err
}
// 返回访问URL和key
protocol := "http"
if m.UseSSL {
protocol = "https"
}
accessURL := protocol + "://" + m.Endpoint + "/" + m.Bucket + "/" + key
return accessURL, key, nil
return m.BucketURL + "/" + info.Key, filePathRes, nil
}
// DeleteFile 删除文件
func (m *Minio) DeleteFile(key string) error {
client, err := m.getClient()
if err != nil {
return err
}
return client.RemoveObject(context.Background(), m.Bucket, key, minio.RemoveObjectOptions{})
ctx, cancel := context.WithTimeout(context.Background(), time.Second*5)
defer cancel()
// Delete the object from MinIO
err := m.Client.RemoveObject(ctx, m.Bucket, key, minio.RemoveObjectOptions{})
return err
}
+145
View File
@@ -29,6 +29,151 @@ type Config struct {
MaxSize int64
}
// OSSConfig 完整的OSS配置
type OSSConfig struct {
OssType string
// Local配置
LocalPath string
LocalStorePath string
// 七牛云配置
QiniuZone string
QiniuBucket string
QiniuImgPath string
QiniuAccessKey string
QiniuSecretKey string
QiniuUseHTTPS bool
QiniuUseCdnDomains bool
// 腾讯云COS配置
TencentBucket string
TencentRegion string
TencentSecretID string
TencentSecretKey string
TencentBaseURL string
TencentPathPrefix string
// 阿里云OSS配置
AliyunEndpoint string
AliyunAccessKeyID string
AliyunAccessKeySecret string
AliyunBucketName string
AliyunBucketURL string
AliyunBasePath string
// 华为云OBS配置
HuaweiPath string
HuaweiBucket string
HuaweiEndpoint string
HuaweiAccessKey string
HuaweiSecretKey string
// AWS S3配置
AwsBucket string
AwsRegion string
AwsEndpoint string
AwsSecretID string
AwsSecretKey string
AwsBaseURL string
AwsPathPrefix string
AwsS3ForcePathStyle bool
AwsDisableSSL bool
// Cloudflare R2配置
CloudflareBucket string
CloudflareBaseURL string
CloudflarePath string
CloudflareAccountID string
CloudflareAccessKeyID string
CloudflareSecretAccessKey string
// MinIO配置
MinioEndpoint string
MinioAccessKeyID string
MinioAccessKeySecret string
MinioBucketName string
MinioUseSSL bool
MinioBasePath string
MinioBucketURL string
}
// NewOSS 根据配置创建OSS实例
func NewOSS(cfg *OSSConfig) (OSS, error) {
switch cfg.OssType {
case "local":
return NewLocal(cfg.LocalStorePath, cfg.LocalPath), nil
case "qiniu":
return NewQiniu(
cfg.QiniuZone,
cfg.QiniuBucket,
cfg.QiniuImgPath,
cfg.QiniuAccessKey,
cfg.QiniuSecretKey,
cfg.QiniuUseHTTPS,
cfg.QiniuUseCdnDomains,
), nil
case "tencent-cos":
return NewTencent(
cfg.TencentBucket,
cfg.TencentRegion,
cfg.TencentSecretID,
cfg.TencentSecretKey,
cfg.TencentBaseURL,
cfg.TencentPathPrefix,
), nil
case "aliyun-oss":
return NewAliyun(
cfg.AliyunEndpoint,
cfg.AliyunAccessKeyID,
cfg.AliyunAccessKeySecret,
cfg.AliyunBucketName,
cfg.AliyunBucketURL,
cfg.AliyunBasePath,
), nil
case "huawei-obs":
return NewHuaweiObs(
cfg.HuaweiPath,
cfg.HuaweiBucket,
cfg.HuaweiEndpoint,
cfg.HuaweiAccessKey,
cfg.HuaweiSecretKey,
), nil
case "aws-s3":
return NewAwsS3(
cfg.AwsBucket,
cfg.AwsRegion,
cfg.AwsEndpoint,
cfg.AwsSecretID,
cfg.AwsSecretKey,
cfg.AwsBaseURL,
cfg.AwsPathPrefix,
cfg.AwsS3ForcePathStyle,
cfg.AwsDisableSSL,
), nil
case "cloudflare-r2":
return NewCloudflareR2(
cfg.CloudflareBucket,
cfg.CloudflareBaseURL,
cfg.CloudflarePath,
cfg.CloudflareAccountID,
cfg.CloudflareAccessKeyID,
cfg.CloudflareSecretAccessKey,
), nil
case "minio":
return NewMinio(
cfg.MinioAccessKeyID,
cfg.MinioAccessKeySecret,
cfg.MinioBucketName,
cfg.MinioEndpoint,
cfg.MinioBasePath,
cfg.MinioUseSSL,
), nil
default:
return NewLocal(cfg.LocalStorePath, cfg.LocalPath), nil
}
}
// GenerateFileName 生成文件名
func GenerateFileName(originalName string) string {
ext := path.Ext(originalName)
+197
View File
@@ -141,3 +141,200 @@ func VariableExistsInBlock(block *ast.BlockStmt, varName string) bool {
})
return exists
}
// SysBaseMenu 菜单结构(用于AST生成)
type SysBaseMenu struct {
ParentId uint
Path string
Name string
Hidden bool
Component string
Sort int
Title string
Icon string
Parameters []SysBaseMenuParameter
MenuBtn []SysBaseMenuBtn
}
// SysBaseMenuParameter 菜单参数
type SysBaseMenuParameter struct {
Type string
Key string
Value string
}
// SysBaseMenuBtn 菜单按钮
type SysBaseMenuBtn struct {
Name string
Desc string
}
// SysApi API结构
type SysApi struct {
Path string
Description string
ApiGroup string
Method string
}
// CreateMenuStructAst 创建菜单结构AST
func CreateMenuStructAst(menus []SysBaseMenu) *[]ast.Expr {
var menuElts []ast.Expr
for i := range menus {
elts := []ast.Expr{
&ast.KeyValueExpr{
Key: &ast.Ident{Name: "ParentId"},
Value: &ast.BasicLit{Kind: token.INT, Value: "0"},
},
&ast.KeyValueExpr{
Key: &ast.Ident{Name: "Path"},
Value: &ast.BasicLit{Kind: token.STRING, Value: fmt.Sprintf("\"%s\"", menus[i].Path)},
},
&ast.KeyValueExpr{
Key: &ast.Ident{Name: "Name"},
Value: &ast.BasicLit{Kind: token.STRING, Value: fmt.Sprintf("\"%s\"", menus[i].Name)},
},
&ast.KeyValueExpr{
Key: &ast.Ident{Name: "Hidden"},
Value: &ast.Ident{Name: "false"},
},
&ast.KeyValueExpr{
Key: &ast.Ident{Name: "Component"},
Value: &ast.BasicLit{Kind: token.STRING, Value: fmt.Sprintf("\"%s\"", menus[i].Component)},
},
&ast.KeyValueExpr{
Key: &ast.Ident{Name: "Sort"},
Value: &ast.BasicLit{Kind: token.INT, Value: fmt.Sprintf("%d", menus[i].Sort)},
},
&ast.KeyValueExpr{
Key: &ast.Ident{Name: "Meta"},
Value: &ast.CompositeLit{
Type: &ast.SelectorExpr{
X: &ast.Ident{Name: "model"},
Sel: &ast.Ident{Name: "Meta"},
},
Elts: []ast.Expr{
&ast.KeyValueExpr{
Key: &ast.Ident{Name: "Title"},
Value: &ast.BasicLit{Kind: token.STRING, Value: fmt.Sprintf("\"%s\"", menus[i].Title)},
},
&ast.KeyValueExpr{
Key: &ast.Ident{Name: "Icon"},
Value: &ast.BasicLit{Kind: token.STRING, Value: fmt.Sprintf("\"%s\"", menus[i].Icon)},
},
},
},
},
}
// 添加菜单参数
if len(menus[i].Parameters) > 0 {
var paramElts []ast.Expr
for _, param := range menus[i].Parameters {
paramElts = append(paramElts, &ast.CompositeLit{
Type: &ast.SelectorExpr{
X: &ast.Ident{Name: "model"},
Sel: &ast.Ident{Name: "SysBaseMenuParameter"},
},
Elts: []ast.Expr{
&ast.KeyValueExpr{
Key: &ast.Ident{Name: "Type"},
Value: &ast.BasicLit{Kind: token.STRING, Value: fmt.Sprintf("\"%s\"", param.Type)},
},
&ast.KeyValueExpr{
Key: &ast.Ident{Name: "Key"},
Value: &ast.BasicLit{Kind: token.STRING, Value: fmt.Sprintf("\"%s\"", param.Key)},
},
&ast.KeyValueExpr{
Key: &ast.Ident{Name: "Value"},
Value: &ast.BasicLit{Kind: token.STRING, Value: fmt.Sprintf("\"%s\"", param.Value)},
},
},
})
}
elts = append(elts, &ast.KeyValueExpr{
Key: &ast.Ident{Name: "Parameters"},
Value: &ast.CompositeLit{
Type: &ast.ArrayType{
Elt: &ast.SelectorExpr{
X: &ast.Ident{Name: "model"},
Sel: &ast.Ident{Name: "SysBaseMenuParameter"},
},
},
Elts: paramElts,
},
})
}
// 添加菜单按钮
if len(menus[i].MenuBtn) > 0 {
var btnElts []ast.Expr
for _, btn := range menus[i].MenuBtn {
btnElts = append(btnElts, &ast.CompositeLit{
Type: &ast.SelectorExpr{
X: &ast.Ident{Name: "model"},
Sel: &ast.Ident{Name: "SysBaseMenuBtn"},
},
Elts: []ast.Expr{
&ast.KeyValueExpr{
Key: &ast.Ident{Name: "Name"},
Value: &ast.BasicLit{Kind: token.STRING, Value: fmt.Sprintf("\"%s\"", btn.Name)},
},
&ast.KeyValueExpr{
Key: &ast.Ident{Name: "Desc"},
Value: &ast.BasicLit{Kind: token.STRING, Value: fmt.Sprintf("\"%s\"", btn.Desc)},
},
},
})
}
elts = append(elts, &ast.KeyValueExpr{
Key: &ast.Ident{Name: "MenuBtn"},
Value: &ast.CompositeLit{
Type: &ast.ArrayType{
Elt: &ast.SelectorExpr{
X: &ast.Ident{Name: "model"},
Sel: &ast.Ident{Name: "SysBaseMenuBtn"},
},
},
Elts: btnElts,
},
})
}
menuElts = append(menuElts, &ast.CompositeLit{
Type: nil,
Elts: elts,
})
}
return &menuElts
}
// CreateApiStructAst 创建API结构AST
func CreateApiStructAst(apis []SysApi) *[]ast.Expr {
var apiElts []ast.Expr
for i := range apis {
elts := []ast.Expr{
&ast.KeyValueExpr{
Key: &ast.Ident{Name: "Path"},
Value: &ast.BasicLit{Kind: token.STRING, Value: fmt.Sprintf("\"%s\"", apis[i].Path)},
},
&ast.KeyValueExpr{
Key: &ast.Ident{Name: "Description"},
Value: &ast.BasicLit{Kind: token.STRING, Value: fmt.Sprintf("\"%s\"", apis[i].Description)},
},
&ast.KeyValueExpr{
Key: &ast.Ident{Name: "ApiGroup"},
Value: &ast.BasicLit{Kind: token.STRING, Value: fmt.Sprintf("\"%s\"", apis[i].ApiGroup)},
},
&ast.KeyValueExpr{
Key: &ast.Ident{Name: "Method"},
Value: &ast.BasicLit{Kind: token.STRING, Value: fmt.Sprintf("\"%s\"", apis[i].Method)},
},
}
apiElts = append(apiElts, &ast.CompositeLit{
Type: nil,
Elts: elts,
})
}
return &apiElts
}
+423
View File
@@ -0,0 +1,423 @@
package ast
import (
"fmt"
"go/ast"
"go/parser"
"go/token"
"log"
)
// AddImport 增加 import 方法
func AddImport(astNode ast.Node, imp string) {
impStr := fmt.Sprintf("\"%s\"", imp)
ast.Inspect(astNode, func(node ast.Node) bool {
if genDecl, ok := node.(*ast.GenDecl); ok {
if genDecl.Tok == token.IMPORT {
for i := range genDecl.Specs {
if impNode, ok := genDecl.Specs[i].(*ast.ImportSpec); ok {
if impNode.Path.Value == impStr {
return false
}
}
}
genDecl.Specs = append(genDecl.Specs, &ast.ImportSpec{
Path: &ast.BasicLit{
Kind: token.STRING,
Value: impStr,
},
})
}
}
return true
})
}
// FindFunction 查询特定function方法
func FindFunction(astNode ast.Node, FunctionName string) *ast.FuncDecl {
var funcDeclP *ast.FuncDecl
ast.Inspect(astNode, func(node ast.Node) bool {
if funcDecl, ok := node.(*ast.FuncDecl); ok {
if funcDecl.Name.String() == FunctionName {
funcDeclP = funcDecl
return false
}
}
return true
})
return funcDeclP
}
// FindArray 查询特定数组方法
func FindArray(astNode ast.Node, identName, selectorExprName string) *ast.CompositeLit {
var assignStmt *ast.CompositeLit
ast.Inspect(astNode, func(n ast.Node) bool {
switch node := n.(type) {
case *ast.AssignStmt:
for _, expr := range node.Rhs {
if exprType, ok := expr.(*ast.CompositeLit); ok {
if arrayType, ok := exprType.Type.(*ast.ArrayType); ok {
sel, ok1 := arrayType.Elt.(*ast.SelectorExpr)
x, ok2 := sel.X.(*ast.Ident)
if ok1 && ok2 && x.Name == identName && sel.Sel.Name == selectorExprName {
assignStmt = exprType
return false
}
}
}
}
}
return true
})
return assignStmt
}
// CheckImport 检查是否存在Import
func CheckImport(file *ast.File, importPath string) bool {
for _, imp := range file.Imports {
path := imp.Path.Value[1 : len(imp.Path.Value)-1]
if path == importPath {
return true
}
}
return false
}
// ClearPosition 清除AST节点位置信息
func ClearPosition(astNode ast.Node) {
ast.Inspect(astNode, func(n ast.Node) bool {
switch node := n.(type) {
case *ast.Ident:
node.NamePos = token.NoPos
case *ast.CallExpr:
node.Lparen = token.NoPos
node.Rparen = token.NoPos
case *ast.BasicLit:
node.ValuePos = token.NoPos
case *ast.SelectorExpr:
node.Sel.NamePos = token.NoPos
case *ast.BinaryExpr:
node.OpPos = token.NoPos
case *ast.UnaryExpr:
node.OpPos = token.NoPos
case *ast.StarExpr:
node.Star = token.NoPos
}
return true
})
}
// CreateStmt 创建语句
func CreateStmt(statement string) *ast.ExprStmt {
expr, err := parser.ParseExpr(statement)
if err != nil {
log.Fatal(err)
}
ClearPosition(expr)
return &ast.ExprStmt{X: expr}
}
// IsBlockStmt 判断是否为块语句
func IsBlockStmt(node ast.Node) bool {
_, ok := node.(*ast.BlockStmt)
return ok
}
// VariableExistsInBlock 检查变量是否存在于块中
func VariableExistsInBlock(block *ast.BlockStmt, varName string) bool {
exists := false
ast.Inspect(block, func(n ast.Node) bool {
switch node := n.(type) {
case *ast.AssignStmt:
for _, expr := range node.Lhs {
if ident, ok := expr.(*ast.Ident); ok && ident.Name == varName {
exists = true
return false
}
}
}
return true
})
return exists
}
// RemoveImport 移除 import
func RemoveImport(file *ast.File, importPath string) {
impStr := fmt.Sprintf("\"%s\"", importPath)
for i := 0; i < len(file.Decls); i++ {
if genDecl, ok := file.Decls[i].(*ast.GenDecl); ok && genDecl.Tok == token.IMPORT {
for j := 0; j < len(genDecl.Specs); j++ {
if impSpec, ok := genDecl.Specs[j].(*ast.ImportSpec); ok {
if impSpec.Path.Value == impStr {
genDecl.Specs = append(genDecl.Specs[:j], genDecl.Specs[j+1:]...)
if len(genDecl.Specs) == 0 {
file.Decls = append(file.Decls[:i], file.Decls[i+1:]...)
}
return
}
}
}
}
}
}
// AddStructField 添加结构体字段
func AddStructField(file *ast.File, structName, fieldName, fieldType string) {
ast.Inspect(file, func(node ast.Node) bool {
if typeSpec, ok := node.(*ast.TypeSpec); ok && typeSpec.Name.Name == structName {
if structType, ok := typeSpec.Type.(*ast.StructType); ok {
// 检查字段是否已存在
for _, field := range structType.Fields.List {
for _, name := range field.Names {
if name.Name == fieldName {
return false
}
}
}
// 添加新字段
newField := &ast.Field{
Names: []*ast.Ident{{Name: fieldName}},
Type: &ast.Ident{Name: fieldType},
}
structType.Fields.List = append(structType.Fields.List, newField)
return false
}
}
return true
})
}
// RemoveStructField 移除结构体字段
func RemoveStructField(file *ast.File, structName, fieldName string) {
ast.Inspect(file, func(node ast.Node) bool {
if typeSpec, ok := node.(*ast.TypeSpec); ok && typeSpec.Name.Name == structName {
if structType, ok := typeSpec.Type.(*ast.StructType); ok {
for i, field := range structType.Fields.List {
for _, name := range field.Names {
if name.Name == fieldName {
structType.Fields.List = append(structType.Fields.List[:i], structType.Fields.List[i+1:]...)
return false
}
}
}
}
}
return true
})
}
// AddFuncCall 在函数体中添加函数调用
func AddFuncCall(file *ast.File, funcName, callExpr string) {
funcDecl := FindFunction(file, funcName)
if funcDecl == nil || funcDecl.Body == nil {
return
}
stmt := CreateStmt(callExpr)
funcDecl.Body.List = append(funcDecl.Body.List, stmt)
}
// RemoveFuncCall 从函数体中移除函数调用
func RemoveFuncCall(file *ast.File, funcName, targetFunc string) {
funcDecl := FindFunction(file, funcName)
if funcDecl == nil || funcDecl.Body == nil {
return
}
for i := 0; i < len(funcDecl.Body.List); i++ {
if exprStmt, ok := funcDecl.Body.List[i].(*ast.ExprStmt); ok {
if callExpr, ok := exprStmt.X.(*ast.CallExpr); ok {
if selExpr, ok := callExpr.Fun.(*ast.SelectorExpr); ok {
if selExpr.Sel.Name == targetFunc {
funcDecl.Body.List = append(funcDecl.Body.List[:i], funcDecl.Body.List[i+1:]...)
return
}
}
}
}
}
}
// AddAutoMigrateModel 添加 AutoMigrate 模型
func AddAutoMigrateModel(file *ast.File, modelExpr string) {
ast.Inspect(file, func(node ast.Node) bool {
if callExpr, ok := node.(*ast.CallExpr); ok {
if selExpr, ok := callExpr.Fun.(*ast.SelectorExpr); ok {
if selExpr.Sel.Name == "AutoMigrate" {
// 检查是否已存在
for _, arg := range callExpr.Args {
if unaryExpr, ok := arg.(*ast.UnaryExpr); ok {
if compLit, ok := unaryExpr.X.(*ast.CompositeLit); ok {
if selType, ok := compLit.Type.(*ast.SelectorExpr); ok {
existingModel := fmt.Sprintf("&%s.%s{}", selType.X.(*ast.Ident).Name, selType.Sel.Name)
if existingModel == modelExpr {
return false
}
}
}
}
}
// 添加新模型
expr, err := parser.ParseExpr(modelExpr)
if err == nil {
ClearPosition(expr)
callExpr.Args = append(callExpr.Args, expr)
}
return false
}
}
}
return true
})
}
// RemoveAutoMigrateModel 移除 AutoMigrate 模型
func RemoveAutoMigrateModel(file *ast.File, modelName string) {
ast.Inspect(file, func(node ast.Node) bool {
if callExpr, ok := node.(*ast.CallExpr); ok {
if selExpr, ok := callExpr.Fun.(*ast.SelectorExpr); ok {
if selExpr.Sel.Name == "AutoMigrate" {
for i := 0; i < len(callExpr.Args); i++ {
if unaryExpr, ok := callExpr.Args[i].(*ast.UnaryExpr); ok {
if compLit, ok := unaryExpr.X.(*ast.CompositeLit); ok {
if selType, ok := compLit.Type.(*ast.SelectorExpr); ok {
existingModel := fmt.Sprintf("%s.%s", selType.X.(*ast.Ident).Name, selType.Sel.Name)
if existingModel == modelName {
callExpr.Args = append(callExpr.Args[:i], callExpr.Args[i+1:]...)
return false
}
}
}
}
}
}
}
}
return true
})
}
// AddGormGenModel 向 GORM Gen 的 g.ApplyBasic 调用中添加新模型
// 用于自动更新 cmd/gen/main.go 文件
func AddGormGenModel(file *ast.File, tableName string) {
ast.Inspect(file, func(node ast.Node) bool {
if callExpr, ok := node.(*ast.CallExpr); ok {
if selExpr, ok := callExpr.Fun.(*ast.SelectorExpr); ok {
// 查找 g.ApplyBasic 调用
if ident, ok := selExpr.X.(*ast.Ident); ok && ident.Name == "g" && selExpr.Sel.Name == "ApplyBasic" {
// 检查是否已存在该表
for _, arg := range callExpr.Args {
if call, ok := arg.(*ast.CallExpr); ok {
if sel, ok := call.Fun.(*ast.SelectorExpr); ok {
if sel.Sel.Name == "GenerateModel" && len(call.Args) > 0 {
if lit, ok := call.Args[0].(*ast.BasicLit); ok {
existingTable := lit.Value[1 : len(lit.Value)-1] // 去掉引号
if existingTable == tableName {
return false // 已存在,不需要添加
}
}
}
}
}
}
// 添加新模型: g.GenerateModel("tableName", fieldOpts...)
newArg := &ast.CallExpr{
Fun: &ast.SelectorExpr{
X: &ast.Ident{Name: "g"},
Sel: &ast.Ident{Name: "GenerateModel"},
},
Args: []ast.Expr{
&ast.BasicLit{Kind: token.STRING, Value: fmt.Sprintf("\"%s\"", tableName)},
&ast.Ident{Name: "fieldOpts..."},
},
}
callExpr.Args = append(callExpr.Args, newArg)
return false
}
}
}
return true
})
}
// RemoveGormGenModel 从 GORM Gen 的 g.ApplyBasic 调用中移除模型
func RemoveGormGenModel(file *ast.File, tableName string) {
ast.Inspect(file, func(node ast.Node) bool {
if callExpr, ok := node.(*ast.CallExpr); ok {
if selExpr, ok := callExpr.Fun.(*ast.SelectorExpr); ok {
if ident, ok := selExpr.X.(*ast.Ident); ok && ident.Name == "g" && selExpr.Sel.Name == "ApplyBasic" {
newArgs := make([]ast.Expr, 0, len(callExpr.Args))
for _, arg := range callExpr.Args {
if call, ok := arg.(*ast.CallExpr); ok {
if sel, ok := call.Fun.(*ast.SelectorExpr); ok {
if sel.Sel.Name == "GenerateModel" && len(call.Args) > 0 {
if lit, ok := call.Args[0].(*ast.BasicLit); ok {
existingTable := lit.Value[1 : len(lit.Value)-1]
if existingTable == tableName {
continue // 跳过要移除的模型
}
}
}
}
}
newArgs = append(newArgs, arg)
}
callExpr.Args = newArgs
return false
}
}
}
return true
})
}
// AddGormGenModelToSlice 向 models 切片中添加新模型
// 用于处理 cmd/gen/main.go 中的 models = append(models, g.GenerateModel(...)) 模式
func AddGormGenModelToSlice(file *ast.File, tableName string) bool {
var found bool
ast.Inspect(file, func(node ast.Node) bool {
// 查找 for 循环中的 models = append(models, ...) 语句
if forStmt, ok := node.(*ast.ForStmt); ok {
ast.Inspect(forStmt.Body, func(n ast.Node) bool {
if assignStmt, ok := n.(*ast.AssignStmt); ok {
for _, rhs := range assignStmt.Rhs {
if callExpr, ok := rhs.(*ast.CallExpr); ok {
if ident, ok := callExpr.Fun.(*ast.Ident); ok && ident.Name == "append" {
// 找到 append 调用,检查是否是 models 切片
if len(callExpr.Args) >= 1 {
if argIdent, ok := callExpr.Args[0].(*ast.Ident); ok && argIdent.Name == "models" {
found = true
return false
}
}
}
}
}
}
return true
})
}
return true
})
return found
}
// CheckGormGenModelExists 检查 GORM Gen 模型是否已存在
func CheckGormGenModelExists(file *ast.File, tableName string) bool {
var exists bool
ast.Inspect(file, func(node ast.Node) bool {
if callExpr, ok := node.(*ast.CallExpr); ok {
if selExpr, ok := callExpr.Fun.(*ast.SelectorExpr); ok {
if selExpr.Sel.Name == "GenerateModel" && len(callExpr.Args) > 0 {
if lit, ok := callExpr.Args[0].(*ast.BasicLit); ok {
existingTable := lit.Value[1 : len(lit.Value)-1]
if existingTable == tableName {
exists = true
return false
}
}
}
}
}
return true
})
return exists
}
+518
View File
@@ -0,0 +1,518 @@
package ast
import (
"go/ast"
"go/parser"
"go/token"
"strings"
"testing"
)
// TestAddImport 测试添加导入
func TestAddImport(t *testing.T) {
testCode := `package main
import "fmt"
func main() {}
`
fset := token.NewFileSet()
file, err := parser.ParseFile(fset, "test.go", testCode, parser.ParseComments)
if err != nil {
t.Fatalf("解析测试代码失败: %v", err)
}
// 添加新导入
AddImport(file, "kra/internal/biz/article")
// 验证导入是否添加
var found bool
ast.Inspect(file, func(node ast.Node) bool {
if genDecl, ok := node.(*ast.GenDecl); ok && genDecl.Tok == token.IMPORT {
for _, spec := range genDecl.Specs {
if impSpec, ok := spec.(*ast.ImportSpec); ok {
if impSpec.Path.Value == "\"kra/internal/biz/article\"" {
found = true
return false
}
}
}
}
return true
})
if !found {
t.Error("应该添加 kra/internal/biz/article 导入")
}
}
// TestAddImport_Duplicate 测试重复添加导入
func TestAddImport_Duplicate(t *testing.T) {
testCode := `package main
import (
"fmt"
"kra/internal/biz/article"
)
func main() {}
`
fset := token.NewFileSet()
file, err := parser.ParseFile(fset, "test.go", testCode, parser.ParseComments)
if err != nil {
t.Fatalf("解析测试代码失败: %v", err)
}
// 统计初始导入数量
initialCount := 0
ast.Inspect(file, func(node ast.Node) bool {
if genDecl, ok := node.(*ast.GenDecl); ok && genDecl.Tok == token.IMPORT {
initialCount = len(genDecl.Specs)
return false
}
return true
})
// 尝试添加已存在的导入
AddImport(file, "kra/internal/biz/article")
// 验证导入数量没有增加
finalCount := 0
ast.Inspect(file, func(node ast.Node) bool {
if genDecl, ok := node.(*ast.GenDecl); ok && genDecl.Tok == token.IMPORT {
finalCount = len(genDecl.Specs)
return false
}
return true
})
if finalCount != initialCount {
t.Errorf("重复添加导入后数量应保持 %d, 实际为 %d", initialCount, finalCount)
}
}
// TestRemoveImport 测试移除导入
func TestRemoveImport(t *testing.T) {
testCode := `package main
import (
"fmt"
"kra/internal/biz/article"
)
func main() {}
`
fset := token.NewFileSet()
file, err := parser.ParseFile(fset, "test.go", testCode, parser.ParseComments)
if err != nil {
t.Fatalf("解析测试代码失败: %v", err)
}
// 移除导入
RemoveImport(file, "kra/internal/biz/article")
// 验证导入是否移除
var found bool
ast.Inspect(file, func(node ast.Node) bool {
if genDecl, ok := node.(*ast.GenDecl); ok && genDecl.Tok == token.IMPORT {
for _, spec := range genDecl.Specs {
if impSpec, ok := spec.(*ast.ImportSpec); ok {
if impSpec.Path.Value == "\"kra/internal/biz/article\"" {
found = true
return false
}
}
}
}
return true
})
if found {
t.Error("应该移除 kra/internal/biz/article 导入")
}
}
// TestAddStructField 测试添加结构体字段
func TestAddStructField(t *testing.T) {
testCode := `package main
type MyStruct struct {
Name string
}
`
fset := token.NewFileSet()
file, err := parser.ParseFile(fset, "test.go", testCode, parser.ParseComments)
if err != nil {
t.Fatalf("解析测试代码失败: %v", err)
}
// 添加新字段
AddStructField(file, "MyStruct", "Age", "int")
// 验证字段是否添加
var found bool
ast.Inspect(file, func(node ast.Node) bool {
if typeSpec, ok := node.(*ast.TypeSpec); ok && typeSpec.Name.Name == "MyStruct" {
if structType, ok := typeSpec.Type.(*ast.StructType); ok {
for _, field := range structType.Fields.List {
for _, name := range field.Names {
if name.Name == "Age" {
found = true
return false
}
}
}
}
}
return true
})
if !found {
t.Error("应该添加 Age 字段")
}
}
// TestAddStructField_Duplicate 测试重复添加结构体字段
func TestAddStructField_Duplicate(t *testing.T) {
testCode := `package main
type MyStruct struct {
Name string
Age int
}
`
fset := token.NewFileSet()
file, err := parser.ParseFile(fset, "test.go", testCode, parser.ParseComments)
if err != nil {
t.Fatalf("解析测试代码失败: %v", err)
}
// 统计初始字段数量
initialCount := 0
ast.Inspect(file, func(node ast.Node) bool {
if typeSpec, ok := node.(*ast.TypeSpec); ok && typeSpec.Name.Name == "MyStruct" {
if structType, ok := typeSpec.Type.(*ast.StructType); ok {
initialCount = len(structType.Fields.List)
return false
}
}
return true
})
// 尝试添加已存在的字段
AddStructField(file, "MyStruct", "Age", "int")
// 验证字段数量没有增加
finalCount := 0
ast.Inspect(file, func(node ast.Node) bool {
if typeSpec, ok := node.(*ast.TypeSpec); ok && typeSpec.Name.Name == "MyStruct" {
if structType, ok := typeSpec.Type.(*ast.StructType); ok {
finalCount = len(structType.Fields.List)
return false
}
}
return true
})
if finalCount != initialCount {
t.Errorf("重复添加字段后数量应保持 %d, 实际为 %d", initialCount, finalCount)
}
}
// TestRemoveStructField 测试移除结构体字段
func TestRemoveStructField(t *testing.T) {
testCode := `package main
type MyStruct struct {
Name string
Age int
}
`
fset := token.NewFileSet()
file, err := parser.ParseFile(fset, "test.go", testCode, parser.ParseComments)
if err != nil {
t.Fatalf("解析测试代码失败: %v", err)
}
// 移除字段
RemoveStructField(file, "MyStruct", "Age")
// 验证字段是否移除
var found bool
ast.Inspect(file, func(node ast.Node) bool {
if typeSpec, ok := node.(*ast.TypeSpec); ok && typeSpec.Name.Name == "MyStruct" {
if structType, ok := typeSpec.Type.(*ast.StructType); ok {
for _, field := range structType.Fields.List {
for _, name := range field.Names {
if name.Name == "Age" {
found = true
return false
}
}
}
}
}
return true
})
if found {
t.Error("应该移除 Age 字段")
}
}
// TestFindFunction 测试查找函数
func TestFindFunction(t *testing.T) {
testCode := `package main
func Hello() {}
func World() {}
`
fset := token.NewFileSet()
file, err := parser.ParseFile(fset, "test.go", testCode, parser.ParseComments)
if err != nil {
t.Fatalf("解析测试代码失败: %v", err)
}
// 查找存在的函数
funcDecl := FindFunction(file, "Hello")
if funcDecl == nil {
t.Error("应该找到 Hello 函数")
}
// 查找不存在的函数
funcDecl = FindFunction(file, "NotExist")
if funcDecl != nil {
t.Error("不应该找到 NotExist 函数")
}
}
// TestCheckImport 测试检查导入是否存在
func TestCheckImport(t *testing.T) {
testCode := `package main
import (
"fmt"
"kra/internal/biz/article"
)
func main() {}
`
fset := token.NewFileSet()
file, err := parser.ParseFile(fset, "test.go", testCode, parser.ParseComments)
if err != nil {
t.Fatalf("解析测试代码失败: %v", err)
}
// 检查存在的导入
if !CheckImport(file, "fmt") {
t.Error("应该找到 fmt 导入")
}
// 检查存在的导入(完整路径)
if !CheckImport(file, "kra/internal/biz/article") {
t.Error("应该找到 kra/internal/biz/article 导入")
}
// 检查不存在的导入
if CheckImport(file, "not/exist") {
t.Error("不应该找到 not/exist 导入")
}
}
// TestAddFuncCall 测试在函数体中添加函数调用
func TestAddFuncCall(t *testing.T) {
testCode := `package main
func initRouter() {
fmt.Println("init")
}
`
fset := token.NewFileSet()
file, err := parser.ParseFile(fset, "test.go", testCode, parser.ParseComments)
if err != nil {
t.Fatalf("解析测试代码失败: %v", err)
}
// 添加函数调用
AddFuncCall(file, "initRouter", "router.Article.InitArticleRouter(group)")
// 验证函数调用是否添加
funcDecl := FindFunction(file, "initRouter")
if funcDecl == nil || funcDecl.Body == nil {
t.Fatal("找不到 initRouter 函数")
}
// 检查函数体中是否包含新的调用
if len(funcDecl.Body.List) < 2 {
t.Error("应该添加新的函数调用")
}
}
// TestRemoveFuncCall 测试从函数体中移除函数调用
func TestRemoveFuncCall(t *testing.T) {
testCode := `package main
func initRouter() {
fmt.Println("init")
router.Article.InitArticleRouter(group)
}
`
fset := token.NewFileSet()
file, err := parser.ParseFile(fset, "test.go", testCode, parser.ParseComments)
if err != nil {
t.Fatalf("解析测试代码失败: %v", err)
}
// 移除函数调用
RemoveFuncCall(file, "initRouter", "InitArticleRouter")
// 验证函数调用是否移除
funcDecl := FindFunction(file, "initRouter")
if funcDecl == nil || funcDecl.Body == nil {
t.Fatal("找不到 initRouter 函数")
}
// 检查函数体中是否还包含该调用
for _, stmt := range funcDecl.Body.List {
if exprStmt, ok := stmt.(*ast.ExprStmt); ok {
if callExpr, ok := exprStmt.X.(*ast.CallExpr); ok {
if selExpr, ok := callExpr.Fun.(*ast.SelectorExpr); ok {
if selExpr.Sel.Name == "InitArticleRouter" {
t.Error("应该移除 InitArticleRouter 调用")
}
}
}
}
}
}
// TestCreateStmt 测试创建语句
func TestCreateStmt(t *testing.T) {
stmt := CreateStmt("fmt.Println(\"hello\")")
if stmt == nil {
t.Error("应该创建语句")
}
if stmt.X == nil {
t.Error("语句表达式不应为空")
}
}
// TestVariableExistsInBlock 测试检查变量是否存在于块中
func TestVariableExistsInBlock(t *testing.T) {
testCode := `package main
func test() {
x := 1
y := 2
}
`
fset := token.NewFileSet()
file, err := parser.ParseFile(fset, "test.go", testCode, parser.ParseComments)
if err != nil {
t.Fatalf("解析测试代码失败: %v", err)
}
funcDecl := FindFunction(file, "test")
if funcDecl == nil || funcDecl.Body == nil {
t.Fatal("找不到 test 函数")
}
// 检查存在的变量
if !VariableExistsInBlock(funcDecl.Body, "x") {
t.Error("应该找到变量 x")
}
// 检查不存在的变量
if VariableExistsInBlock(funcDecl.Body, "z") {
t.Error("不应该找到变量 z")
}
}
// TestBase_RelativePath 测试相对路径转换
func TestBase_RelativePath(t *testing.T) {
base := &Base{
AutoCodeRoot: "/home/user/project",
AutoCodeServer: "server",
}
testCases := []struct {
input string
expected string
}{
{"/home/user/project/server/internal/biz/article.go", "internal/biz/article.go"},
{"/other/path/file.go", "/other/path/file.go"},
}
for _, tc := range testCases {
result := base.RelativePath(tc.input)
// 由于路径分隔符可能不同,只检查是否包含关键部分
if !strings.Contains(result, "article.go") && tc.input != "/other/path/file.go" {
t.Errorf("RelativePath(%s) 结果不正确: %s", tc.input, result)
}
}
}
// TestAddAutoMigrateModel 测试添加 AutoMigrate 模型
func TestAddAutoMigrateModel(t *testing.T) {
testCode := `package main
import "gorm.io/gorm"
func migrate(db *gorm.DB) {
db.AutoMigrate(&User{})
}
`
fset := token.NewFileSet()
file, err := parser.ParseFile(fset, "test.go", testCode, parser.ParseComments)
if err != nil {
t.Fatalf("解析测试代码失败: %v", err)
}
// 添加新模型
AddAutoMigrateModel(file, "&article.Article{}")
// 验证模型是否添加
var found bool
ast.Inspect(file, func(node ast.Node) bool {
if callExpr, ok := node.(*ast.CallExpr); ok {
if selExpr, ok := callExpr.Fun.(*ast.SelectorExpr); ok {
if selExpr.Sel.Name == "AutoMigrate" {
// 检查参数数量是否增加
if len(callExpr.Args) > 1 {
found = true
return false
}
}
}
}
return true
})
if !found {
t.Error("应该添加新的 AutoMigrate 模型")
}
}
// TestCheckGormGenModelExists 测试检查 GORM Gen 模型是否存在
func TestCheckGormGenModelExists(t *testing.T) {
testCode := `package main
func main() {
g.ApplyBasic(
g.GenerateModel("users"),
g.GenerateModel("articles"),
)
}
`
fset := token.NewFileSet()
file, err := parser.ParseFile(fset, "test.go", testCode, parser.ParseComments)
if err != nil {
t.Fatalf("解析测试代码失败: %v", err)
}
// 检查存在的模型
if !CheckGormGenModelExists(file, "users") {
t.Error("应该找到 users 模型")
}
// 检查不存在的模型
if CheckGormGenModelExists(file, "orders") {
t.Error("不应该找到 orders 模型")
}
}
+217
View File
@@ -0,0 +1,217 @@
package ast
import (
"bytes"
"go/ast"
"go/format"
"go/parser"
"go/token"
"log"
"os"
"strconv"
"strings"
"golang.org/x/text/cases"
"golang.org/x/text/language"
)
// Visitor AST访问器
type Visitor struct {
ImportCode string
StructName string
PackageName string
GroupName string
}
func (vi *Visitor) Visit(node ast.Node) ast.Visitor {
switch n := node.(type) {
case *ast.GenDecl:
if n.Tok == token.IMPORT && vi.ImportCode != "" {
vi.addImport(n)
return nil
}
if n.Tok == token.TYPE && vi.StructName != "" && vi.PackageName != "" && vi.GroupName != "" {
vi.addStruct(n)
return nil
}
case *ast.FuncDecl:
if n.Name.Name == "Routers" {
vi.addFuncBodyVar(n)
return nil
}
}
return vi
}
func (vi *Visitor) addStruct(genDecl *ast.GenDecl) ast.Visitor {
for i := range genDecl.Specs {
switch n := genDecl.Specs[i].(type) {
case *ast.TypeSpec:
if strings.Index(n.Name.Name, "Group") > -1 {
switch t := n.Type.(type) {
case *ast.StructType:
f := &ast.Field{
Names: []*ast.Ident{
{
Name: vi.StructName,
Obj: &ast.Object{
Kind: ast.Var,
Name: vi.StructName,
},
},
},
Type: &ast.SelectorExpr{
X: &ast.Ident{
Name: vi.PackageName,
},
Sel: &ast.Ident{
Name: vi.GroupName,
},
},
}
t.Fields.List = append(t.Fields.List, f)
}
}
}
}
return vi
}
func (vi *Visitor) addImport(genDecl *ast.GenDecl) ast.Visitor {
hasImported := false
for _, v := range genDecl.Specs {
importSpec := v.(*ast.ImportSpec)
if importSpec.Path.Value == strconv.Quote(vi.ImportCode) {
hasImported = true
}
}
if !hasImported {
genDecl.Specs = append(genDecl.Specs, &ast.ImportSpec{
Path: &ast.BasicLit{
Kind: token.STRING,
Value: strconv.Quote(vi.ImportCode),
},
})
}
return vi
}
func (vi *Visitor) addFuncBodyVar(funDecl *ast.FuncDecl) ast.Visitor {
hasVar := false
for _, v := range funDecl.Body.List {
switch varSpec := v.(type) {
case *ast.AssignStmt:
for i := range varSpec.Lhs {
switch nn := varSpec.Lhs[i].(type) {
case *ast.Ident:
if nn.Name == vi.PackageName+"Router" {
hasVar = true
}
}
}
}
}
if !hasVar {
assignStmt := &ast.AssignStmt{
Lhs: []ast.Expr{
&ast.Ident{
Name: vi.PackageName + "Router",
Obj: &ast.Object{
Kind: ast.Var,
Name: vi.PackageName + "Router",
},
},
},
Tok: token.DEFINE,
Rhs: []ast.Expr{
&ast.SelectorExpr{
X: &ast.SelectorExpr{
X: &ast.Ident{
Name: "router",
},
Sel: &ast.Ident{
Name: "RouterGroupApp",
},
},
Sel: &ast.Ident{
Name: cases.Title(language.English).String(vi.PackageName),
},
},
},
}
funDecl.Body.List = append(funDecl.Body.List, funDecl.Body.List[1])
index := 1
copy(funDecl.Body.List[index+1:], funDecl.Body.List[index:])
funDecl.Body.List[index] = assignStmt
}
return vi
}
// ImportReference 导入引用
func ImportReference(filepath, importCode, structName, packageName, groupName string) error {
fSet := token.NewFileSet()
fParser, err := parser.ParseFile(fSet, filepath, nil, parser.ParseComments)
if err != nil {
return err
}
importCode = strings.TrimSpace(importCode)
v := &Visitor{
ImportCode: importCode,
StructName: structName,
PackageName: packageName,
GroupName: groupName,
}
if importCode == "" {
ast.Print(fSet, fParser)
}
ast.Walk(v, fParser)
var output []byte
buffer := bytes.NewBuffer(output)
err = format.Node(buffer, fSet, fParser)
if err != nil {
log.Fatal(err)
}
return os.WriteFile(filepath, buffer.Bytes(), 0o600)
}
// ImportForAutoEnter 自动导入
func ImportForAutoEnter(path string, funcName string, code string) {
src, err := os.ReadFile(path)
if err != nil {
return
}
fileSet := token.NewFileSet()
astFile, err := parser.ParseFile(fileSet, "", src, 0)
if err != nil {
return
}
ast.Inspect(astFile, func(node ast.Node) bool {
if typeSpec, ok := node.(*ast.TypeSpec); ok {
if typeSpec.Name.Name == funcName {
if st, ok := typeSpec.Type.(*ast.StructType); ok {
for i := range st.Fields.List {
if t, ok := st.Fields.List[i].Type.(*ast.Ident); ok {
if t.Name == code {
return false
}
}
}
sn := &ast.Field{
Type: &ast.Ident{Name: code},
}
st.Fields.List = append(st.Fields.List, sn)
}
}
}
return true
})
var out []byte
bf := bytes.NewBuffer(out)
err = format.Node(bf, fileSet, astFile)
if err != nil {
return
}
_ = os.WriteFile(path, bf.Bytes(), 0666)
}
+140
View File
@@ -0,0 +1,140 @@
package ast
import (
"bytes"
"fmt"
"go/ast"
"go/parser"
"go/printer"
"go/token"
"os"
)
// AddRegisterTablesAst 自动为 gorm.go 注册一个自动迁移
func AddRegisterTablesAst(path, funcName, pk, varName, dbName, model, modelImportPath string) {
src, err := os.ReadFile(path)
if err != nil {
fmt.Println(err)
return
}
fileSet := token.NewFileSet()
astFile, err := parser.ParseFile(fileSet, "", src, 0)
if err != nil {
fmt.Println(err)
return
}
AddImport(astFile, modelImportPath)
FuncNode := FindFunction(astFile, funcName)
if FuncNode != nil {
addDBVar(FuncNode.Body, varName, dbName)
addAutoMigrate(FuncNode.Body, varName, pk, model)
}
var out []byte
bf := bytes.NewBuffer(out)
printer.Fprint(bf, fileSet, astFile)
os.WriteFile(path, bf.Bytes(), 0666)
}
// addDBVar 增加一个 db库变量
func addDBVar(astBody *ast.BlockStmt, varName, dbName string) {
if dbName == "" {
return
}
dbStr := fmt.Sprintf("\"%s\"", dbName)
for i := range astBody.List {
if assignStmt, ok := astBody.List[i].(*ast.AssignStmt); ok {
if ident, ok := assignStmt.Lhs[0].(*ast.Ident); ok {
if ident.Name == varName {
return
}
}
}
}
assignNode := &ast.AssignStmt{
Lhs: []ast.Expr{
&ast.Ident{Name: varName},
},
Tok: token.DEFINE,
Rhs: []ast.Expr{
&ast.CallExpr{
Fun: &ast.SelectorExpr{
X: &ast.Ident{Name: "global"},
Sel: &ast.Ident{Name: "GetGlobalDBByDBName"},
},
Args: []ast.Expr{
&ast.BasicLit{Kind: token.STRING, Value: dbStr},
},
},
},
}
astBody.List = append([]ast.Stmt{assignNode}, astBody.List...)
}
// addAutoMigrate 为db库变量增加 AutoMigrate 方法
func addAutoMigrate(astBody *ast.BlockStmt, dbname string, pk string, model string) {
if dbname == "" {
dbname = "db"
}
flag := true
ast.Inspect(astBody, func(node ast.Node) bool {
switch n := node.(type) {
case *ast.CallExpr:
if s, ok := n.Fun.(*ast.SelectorExpr); ok {
if x, ok := s.X.(*ast.Ident); ok {
if s.Sel.Name == "AutoMigrate" && x.Name == dbname {
flag = false
if !NeedAppendModel(n, pk, model) {
return false
}
n.Args = append(n.Args, &ast.CompositeLit{
Type: &ast.SelectorExpr{
X: &ast.Ident{Name: pk},
Sel: &ast.Ident{Name: model},
},
})
return false
}
}
}
}
return true
})
if flag {
exprStmt := &ast.ExprStmt{
X: &ast.CallExpr{
Fun: &ast.SelectorExpr{
X: &ast.Ident{Name: dbname},
Sel: &ast.Ident{Name: "AutoMigrate"},
},
Args: []ast.Expr{
&ast.CompositeLit{
Type: &ast.SelectorExpr{
X: &ast.Ident{Name: pk},
Sel: &ast.Ident{Name: model},
},
},
},
}}
astBody.List = append(astBody.List, exprStmt)
}
}
// NeedAppendModel 为automigrate增加实参
func NeedAppendModel(callNode ast.Node, pk string, model string) bool {
flag := true
ast.Inspect(callNode, func(node ast.Node) bool {
switch n := node.(type) {
case *ast.SelectorExpr:
if x, ok := n.X.(*ast.Ident); ok {
if n.Sel.Name == model && x.Name == pk {
flag = false
return false
}
}
}
return true
})
return flag
}
+96
View File
@@ -0,0 +1,96 @@
package ast
import (
"go/ast"
"go/token"
"io"
"strings"
)
// Import 导入操作
type Import struct {
Base
ImportPath string
}
// NewImport 创建导入操作
func NewImport(importPath string) *Import {
return &Import{ImportPath: importPath}
}
func (a *Import) Parse(filename string, writer io.Writer) (file *ast.File, err error) {
return a.Base.Parse(filename, writer)
}
func (a *Import) Rollback(file *ast.File) error {
if a.ImportPath == "" {
return nil
}
for i := 0; i < len(file.Decls); i++ {
v1, o1 := file.Decls[i].(*ast.GenDecl)
if o1 {
if v1.Tok != token.IMPORT {
break
}
for j := 0; j < len(v1.Specs); j++ {
v2, o2 := v1.Specs[j].(*ast.ImportSpec)
if o2 && strings.HasSuffix(a.ImportPath, v2.Path.Value) {
v1.Specs = append(v1.Specs[:j], v1.Specs[j+1:]...)
if len(v1.Specs) == 0 {
file.Decls = append(file.Decls[:i], file.Decls[i+1:]...)
}
break
}
}
}
}
return nil
}
func (a *Import) Injection(file *ast.File) error {
if a.ImportPath == "" {
return nil
}
var has bool
for i := 0; i < len(file.Decls); i++ {
v1, o1 := file.Decls[i].(*ast.GenDecl)
if o1 {
if v1.Tok != token.IMPORT {
break
}
for j := 0; j < len(v1.Specs); j++ {
v2, o2 := v1.Specs[j].(*ast.ImportSpec)
if o2 && strings.HasSuffix(a.ImportPath, v2.Path.Value) {
has = true
break
}
}
if !has {
spec := &ast.ImportSpec{
Path: &ast.BasicLit{Kind: token.STRING, Value: a.ImportPath},
}
v1.Specs = append(v1.Specs, spec)
return nil
}
}
}
if !has {
decls := file.Decls
file.Decls = make([]ast.Decl, 0, len(file.Decls)+1)
decl := &ast.GenDecl{
Tok: token.IMPORT,
Specs: []ast.Spec{
&ast.ImportSpec{
Path: &ast.BasicLit{Kind: token.STRING, Value: a.ImportPath},
},
},
}
file.Decls = append(file.Decls, decl)
file.Decls = append(file.Decls, decls...)
}
return nil
}
func (a *Import) Format(filename string, writer io.Writer, file *ast.File) error {
return a.Base.Format(filename, writer, file)
}
+88
View File
@@ -0,0 +1,88 @@
package ast
import (
"go/ast"
"go/format"
"go/parser"
"go/token"
"io"
"os"
"path"
"path/filepath"
"strings"
"github.com/pkg/errors"
)
// Ast AST操作接口
type Ast interface {
Parse(filename string, writer io.Writer) (file *ast.File, err error)
Rollback(file *ast.File) error
Injection(file *ast.File) error
Format(filename string, writer io.Writer, file *ast.File) error
}
// Base 基础AST操作
type Base struct {
AutoCodeRoot string
AutoCodeServer string
}
func (a *Base) Parse(filename string, writer io.Writer) (file *ast.File, err error) {
fileSet := token.NewFileSet()
if writer != nil {
file, err = parser.ParseFile(fileSet, filename, nil, parser.ParseComments)
} else {
file, err = parser.ParseFile(fileSet, filename, writer, parser.ParseComments)
}
if err != nil {
return nil, errors.Wrapf(err, "[filepath:%s]打开/解析文件失败!", filename)
}
return file, nil
}
func (a *Base) Rollback(file *ast.File) error {
return nil
}
func (a *Base) Injection(file *ast.File) error {
return nil
}
func (a *Base) Format(filename string, writer io.Writer, file *ast.File) error {
fileSet := token.NewFileSet()
if writer == nil {
open, err := os.OpenFile(filename, os.O_WRONLY|os.O_TRUNC, 0666)
defer open.Close()
if err != nil {
return errors.Wrapf(err, "[filepath:%s]打开文件失败!", filename)
}
writer = open
}
err := format.Node(writer, fileSet, file)
if err != nil {
return errors.Wrapf(err, "[filepath:%s]注入失败!", filename)
}
return nil
}
// RelativePath 绝对路径转相对路径
func (a *Base) RelativePath(filePath string) string {
server := filepath.Join(a.AutoCodeRoot, a.AutoCodeServer)
hasServer := strings.Index(filePath, server)
if hasServer != -1 {
filePath = strings.TrimPrefix(filePath, server)
keys := strings.Split(filePath, string(filepath.Separator))
filePath = path.Join(keys...)
}
return filePath
}
// AbsolutePath 相对路径转绝对路径
func (a *Base) AbsolutePath(filePath string) string {
server := filepath.Join(a.AutoCodeRoot, a.AutoCodeServer)
keys := strings.Split(filePath, "/")
filePath = filepath.Join(keys...)
filePath = filepath.Join(server, filePath)
return filePath
}
+135
View File
@@ -0,0 +1,135 @@
package ast
import (
"bytes"
"fmt"
"go/ast"
"go/parser"
"go/printer"
"go/token"
"os"
"strings"
)
// AppendNodeToList 向语句列表中插入节点
func AppendNodeToList(stmts []ast.Stmt, stmt ast.Stmt, index int) []ast.Stmt {
return append(stmts[:index], append([]ast.Stmt{stmt}, stmts[index:]...)...)
}
// AddRouterCode 添加路由代码
func AddRouterCode(path, funcName, pk, model string) {
src, err := os.ReadFile(path)
if err != nil {
fmt.Println(err)
return
}
fileSet := token.NewFileSet()
astFile, err := parser.ParseFile(fileSet, "", src, parser.ParseComments)
if err != nil {
fmt.Println(err)
return
}
FuncNode := FindFunction(astFile, funcName)
if FuncNode == nil {
return
}
pkName := strings.ToUpper(pk[:1]) + pk[1:]
routerName := fmt.Sprintf("%sRouter", pk)
modelName := fmt.Sprintf("Init%sRouter", model)
var bloctPre *ast.BlockStmt
for i := len(FuncNode.Body.List) - 1; i >= 0; i-- {
if block, ok := FuncNode.Body.List[i].(*ast.BlockStmt); ok {
bloctPre = block
}
}
if ok, b := needAppendRouter(FuncNode, pk); ok {
routerNode := &ast.BlockStmt{
List: []ast.Stmt{
&ast.AssignStmt{
Lhs: []ast.Expr{
&ast.Ident{Name: routerName},
},
Tok: token.DEFINE,
Rhs: []ast.Expr{
&ast.SelectorExpr{
X: &ast.SelectorExpr{
X: &ast.Ident{Name: "router"},
Sel: &ast.Ident{Name: "RouterGroupApp"},
},
Sel: &ast.Ident{Name: pkName},
},
},
},
},
}
FuncNode.Body.List = AppendNodeToList(FuncNode.Body.List, routerNode, len(FuncNode.Body.List)-1)
bloctPre = routerNode
} else {
bloctPre = b
}
if needAppendInit(FuncNode, routerName, modelName) {
bloctPre.List = append(bloctPre.List,
&ast.ExprStmt{
X: &ast.CallExpr{
Fun: &ast.SelectorExpr{
X: &ast.Ident{Name: routerName},
Sel: &ast.Ident{Name: modelName},
},
Args: []ast.Expr{
&ast.Ident{Name: "privateGroup"},
&ast.Ident{Name: "publicGroup"},
},
},
})
}
var out []byte
bf := bytes.NewBuffer(out)
printer.Fprint(bf, fileSet, astFile)
os.WriteFile(path, bf.Bytes(), 0666)
}
func needAppendRouter(funcNode ast.Node, pk string) (bool, *ast.BlockStmt) {
flag := true
var block *ast.BlockStmt
ast.Inspect(funcNode, func(node ast.Node) bool {
switch n := node.(type) {
case *ast.BlockStmt:
for i := range n.List {
if assignNode, ok := n.List[i].(*ast.AssignStmt); ok {
if identNode, ok := assignNode.Lhs[0].(*ast.Ident); ok {
if identNode.Name == fmt.Sprintf("%sRouter", pk) {
flag = false
block = n
return false
}
}
}
}
}
return true
})
return flag, block
}
func needAppendInit(funcNode ast.Node, routerName string, modelName string) bool {
flag := true
ast.Inspect(funcNode, func(node ast.Node) bool {
switch n := funcNode.(type) {
case *ast.CallExpr:
if selectNode, ok := n.Fun.(*ast.SelectorExpr); ok {
x, xok := selectNode.X.(*ast.Ident)
if xok && x.Name == routerName && selectNode.Sel.Name == modelName {
flag = false
return false
}
}
}
return true
})
return flag
}
+255
View File
@@ -0,0 +1,255 @@
package ast
import (
"fmt"
"go/ast"
"go/token"
)
// Type AST类型
type Type string
func (r Type) String() string {
return string(r)
}
func (r Type) Group() string {
switch r {
case TypePackageApiEnter:
return "ApiGroup"
case TypePackageRouterEnter:
return "RouterGroup"
case TypePackageServiceEnter:
return "ServiceGroup"
case TypePackageApiModuleEnter:
return "ApiGroup"
case TypePackageRouterModuleEnter:
return "RouterGroup"
case TypePackageServiceModuleEnter:
return "ServiceGroup"
case TypePluginApiEnter:
return "api"
case TypePluginRouterEnter:
return "router"
case TypePluginServiceEnter:
return "service"
default:
return ""
}
}
const (
TypePackageApiEnter = "PackageApiEnter"
TypePackageRouterEnter = "PackageRouterEnter"
TypePackageServiceEnter = "PackageServiceEnter"
TypePackageApiModuleEnter = "PackageApiModuleEnter"
TypePackageRouterModuleEnter = "PackageRouterModuleEnter"
TypePackageServiceModuleEnter = "PackageServiceModuleEnter"
TypePackageInitializeGorm = "PackageInitializeGorm"
TypePackageInitializeRouter = "PackageInitializeRouter"
TypePluginGen = "PluginGen"
TypePluginApiEnter = "PluginApiEnter"
TypePluginInitializeV1 = "PluginInitializeV1"
TypePluginInitializeV2 = "PluginInitializeV2"
TypePluginRouterEnter = "PluginRouterEnter"
TypePluginServiceEnter = "PluginServiceEnter"
TypePluginInitializeApi = "PluginInitializeApi"
TypePluginInitializeGorm = "PluginInitializeGorm"
TypePluginInitializeMenu = "PluginInitializeMenu"
TypePluginInitializeRouter = "PluginInitializeRouter"
)
// SysBaseMenu 菜单结构(用于AST生成)
type SysBaseMenu struct {
ParentId uint
Path string
Name string
Hidden bool
Component string
Sort int
Title string
Icon string
Parameters []SysBaseMenuParameter
MenuBtn []SysBaseMenuBtn
}
// SysBaseMenuParameter 菜单参数
type SysBaseMenuParameter struct {
Type string
Key string
Value string
}
// SysBaseMenuBtn 菜单按钮
type SysBaseMenuBtn struct {
Name string
Desc string
}
// SysApi API结构
type SysApi struct {
Path string
Description string
ApiGroup string
Method string
}
// CreateMenuStructAst 创建菜单结构AST
func CreateMenuStructAst(menus []SysBaseMenu) *[]ast.Expr {
var menuElts []ast.Expr
for i := range menus {
elts := []ast.Expr{
&ast.KeyValueExpr{
Key: &ast.Ident{Name: "ParentId"},
Value: &ast.BasicLit{Kind: token.INT, Value: "0"},
},
&ast.KeyValueExpr{
Key: &ast.Ident{Name: "Path"},
Value: &ast.BasicLit{Kind: token.STRING, Value: fmt.Sprintf("\"%s\"", menus[i].Path)},
},
&ast.KeyValueExpr{
Key: &ast.Ident{Name: "Name"},
Value: &ast.BasicLit{Kind: token.STRING, Value: fmt.Sprintf("\"%s\"", menus[i].Name)},
},
&ast.KeyValueExpr{
Key: &ast.Ident{Name: "Hidden"},
Value: &ast.Ident{Name: "false"},
},
&ast.KeyValueExpr{
Key: &ast.Ident{Name: "Component"},
Value: &ast.BasicLit{Kind: token.STRING, Value: fmt.Sprintf("\"%s\"", menus[i].Component)},
},
&ast.KeyValueExpr{
Key: &ast.Ident{Name: "Sort"},
Value: &ast.BasicLit{Kind: token.INT, Value: fmt.Sprintf("%d", menus[i].Sort)},
},
&ast.KeyValueExpr{
Key: &ast.Ident{Name: "Meta"},
Value: &ast.CompositeLit{
Type: &ast.SelectorExpr{
X: &ast.Ident{Name: "model"},
Sel: &ast.Ident{Name: "Meta"},
},
Elts: []ast.Expr{
&ast.KeyValueExpr{
Key: &ast.Ident{Name: "Title"},
Value: &ast.BasicLit{Kind: token.STRING, Value: fmt.Sprintf("\"%s\"", menus[i].Title)},
},
&ast.KeyValueExpr{
Key: &ast.Ident{Name: "Icon"},
Value: &ast.BasicLit{Kind: token.STRING, Value: fmt.Sprintf("\"%s\"", menus[i].Icon)},
},
},
},
},
}
if len(menus[i].Parameters) > 0 {
var paramElts []ast.Expr
for _, param := range menus[i].Parameters {
paramElts = append(paramElts, &ast.CompositeLit{
Type: &ast.SelectorExpr{
X: &ast.Ident{Name: "model"},
Sel: &ast.Ident{Name: "SysBaseMenuParameter"},
},
Elts: []ast.Expr{
&ast.KeyValueExpr{
Key: &ast.Ident{Name: "Type"},
Value: &ast.BasicLit{Kind: token.STRING, Value: fmt.Sprintf("\"%s\"", param.Type)},
},
&ast.KeyValueExpr{
Key: &ast.Ident{Name: "Key"},
Value: &ast.BasicLit{Kind: token.STRING, Value: fmt.Sprintf("\"%s\"", param.Key)},
},
&ast.KeyValueExpr{
Key: &ast.Ident{Name: "Value"},
Value: &ast.BasicLit{Kind: token.STRING, Value: fmt.Sprintf("\"%s\"", param.Value)},
},
},
})
}
elts = append(elts, &ast.KeyValueExpr{
Key: &ast.Ident{Name: "Parameters"},
Value: &ast.CompositeLit{
Type: &ast.ArrayType{
Elt: &ast.SelectorExpr{
X: &ast.Ident{Name: "model"},
Sel: &ast.Ident{Name: "SysBaseMenuParameter"},
},
},
Elts: paramElts,
},
})
}
if len(menus[i].MenuBtn) > 0 {
var btnElts []ast.Expr
for _, btn := range menus[i].MenuBtn {
btnElts = append(btnElts, &ast.CompositeLit{
Type: &ast.SelectorExpr{
X: &ast.Ident{Name: "model"},
Sel: &ast.Ident{Name: "SysBaseMenuBtn"},
},
Elts: []ast.Expr{
&ast.KeyValueExpr{
Key: &ast.Ident{Name: "Name"},
Value: &ast.BasicLit{Kind: token.STRING, Value: fmt.Sprintf("\"%s\"", btn.Name)},
},
&ast.KeyValueExpr{
Key: &ast.Ident{Name: "Desc"},
Value: &ast.BasicLit{Kind: token.STRING, Value: fmt.Sprintf("\"%s\"", btn.Desc)},
},
},
})
}
elts = append(elts, &ast.KeyValueExpr{
Key: &ast.Ident{Name: "MenuBtn"},
Value: &ast.CompositeLit{
Type: &ast.ArrayType{
Elt: &ast.SelectorExpr{
X: &ast.Ident{Name: "model"},
Sel: &ast.Ident{Name: "SysBaseMenuBtn"},
},
},
Elts: btnElts,
},
})
}
menuElts = append(menuElts, &ast.CompositeLit{
Type: nil,
Elts: elts,
})
}
return &menuElts
}
// CreateApiStructAst 创建API结构AST
func CreateApiStructAst(apis []SysApi) *[]ast.Expr {
var apiElts []ast.Expr
for i := range apis {
elts := []ast.Expr{
&ast.KeyValueExpr{
Key: &ast.Ident{Name: "Path"},
Value: &ast.BasicLit{Kind: token.STRING, Value: fmt.Sprintf("\"%s\"", apis[i].Path)},
},
&ast.KeyValueExpr{
Key: &ast.Ident{Name: "Description"},
Value: &ast.BasicLit{Kind: token.STRING, Value: fmt.Sprintf("\"%s\"", apis[i].Description)},
},
&ast.KeyValueExpr{
Key: &ast.Ident{Name: "ApiGroup"},
Value: &ast.BasicLit{Kind: token.STRING, Value: fmt.Sprintf("\"%s\"", apis[i].ApiGroup)},
},
&ast.KeyValueExpr{
Key: &ast.Ident{Name: "Method"},
Value: &ast.BasicLit{Kind: token.STRING, Value: fmt.Sprintf("\"%s\"", apis[i].Method)},
},
}
apiElts = append(apiElts, &ast.CompositeLit{
Type: nil,
Elts: elts,
})
}
return &apiElts
}
+618
View File
@@ -0,0 +1,618 @@
package autocode
import (
"fmt"
"go/ast"
"go/format"
"go/parser"
"go/token"
"io"
"os"
"github.com/pkg/errors"
)
// WireProviderAst Wire Provider AST 注入结构体
// 用于向各层的 ProviderSet 中注入新的 Provider
type WireProviderAst struct {
Layer string // 层名称: biz, data, handler, service, router
PackageName string // 包名
ProviderName string // Provider 名称,如 NewArticleUsecase
Path string // 目标文件路径
}
// NewWireProviderAst 创建 Wire Provider AST 注入器
func NewWireProviderAst(layer, packageName, providerName, path string) *WireProviderAst {
return &WireProviderAst{
Layer: layer,
PackageName: packageName,
ProviderName: providerName,
Path: path,
}
}
// Parse 解析文件
func (a *WireProviderAst) Parse(filename string, writer io.Writer) (*ast.File, error) {
if filename == "" {
filename = a.Path
}
if filename == "" {
return nil, fmt.Errorf("文件路径为空")
}
fset := token.NewFileSet()
file, err := parser.ParseFile(fset, filename, nil, parser.ParseComments)
if err != nil {
return nil, fmt.Errorf("解析文件失败 %s: %w", filename, err)
}
return file, nil
}
// Rollback 回滚注入
func (a *WireProviderAst) Rollback(file *ast.File) error {
return removeWireProvider(file, a.ProviderName)
}
// Injection 注入代码
func (a *WireProviderAst) Injection(file *ast.File) error {
return addWireProvider(file, a.ProviderName)
}
// Format 格式化输出
func (a *WireProviderAst) Format(filename string, writer io.Writer, file *ast.File) error {
if filename == "" {
filename = a.Path
}
fset := token.NewFileSet()
if writer == nil {
f, err := os.OpenFile(filename, os.O_WRONLY|os.O_TRUNC, 0666)
if err != nil {
return errors.Wrapf(err, "打开文件失败: %s", filename)
}
defer f.Close()
writer = f
}
return format.Node(writer, fset, file)
}
// addWireProvider 向 wire.NewSet 中添加 Provider
func addWireProvider(file *ast.File, providerName string) error {
var found bool
ast.Inspect(file, func(node ast.Node) bool {
// 查找 var ProviderSet = wire.NewSet(...) 声明
genDecl, ok := node.(*ast.GenDecl)
if !ok || genDecl.Tok != token.VAR {
return true
}
for _, spec := range genDecl.Specs {
valueSpec, ok := spec.(*ast.ValueSpec)
if !ok {
continue
}
// 检查是否是 ProviderSet 变量
for i, name := range valueSpec.Names {
if name.Name != "ProviderSet" {
continue
}
if i >= len(valueSpec.Values) {
continue
}
// 查找 wire.NewSet 调用
callExpr, ok := valueSpec.Values[i].(*ast.CallExpr)
if !ok {
continue
}
selExpr, ok := callExpr.Fun.(*ast.SelectorExpr)
if !ok {
continue
}
ident, ok := selExpr.X.(*ast.Ident)
if !ok || ident.Name != "wire" || selExpr.Sel.Name != "NewSet" {
continue
}
// 检查 Provider 是否已存在
for _, arg := range callExpr.Args {
if argIdent, ok := arg.(*ast.Ident); ok {
if argIdent.Name == providerName {
found = true
return false // 已存在,不需要添加
}
}
}
// 添加新的 Provider
newArg := &ast.Ident{Name: providerName}
callExpr.Args = append(callExpr.Args, newArg)
found = true
return false
}
}
return true
})
if !found {
return fmt.Errorf("未找到 ProviderSet 变量或 wire.NewSet 调用")
}
return nil
}
// removeWireProvider 从 wire.NewSet 中移除 Provider
func removeWireProvider(file *ast.File, providerName string) error {
ast.Inspect(file, func(node ast.Node) bool {
genDecl, ok := node.(*ast.GenDecl)
if !ok || genDecl.Tok != token.VAR {
return true
}
for _, spec := range genDecl.Specs {
valueSpec, ok := spec.(*ast.ValueSpec)
if !ok {
continue
}
for i, name := range valueSpec.Names {
if name.Name != "ProviderSet" {
continue
}
if i >= len(valueSpec.Values) {
continue
}
callExpr, ok := valueSpec.Values[i].(*ast.CallExpr)
if !ok {
continue
}
selExpr, ok := callExpr.Fun.(*ast.SelectorExpr)
if !ok {
continue
}
ident, ok := selExpr.X.(*ast.Ident)
if !ok || ident.Name != "wire" || selExpr.Sel.Name != "NewSet" {
continue
}
// 移除指定的 Provider
newArgs := make([]ast.Expr, 0, len(callExpr.Args))
for _, arg := range callExpr.Args {
if argIdent, ok := arg.(*ast.Ident); ok {
if argIdent.Name == providerName {
continue // 跳过要移除的 Provider
}
}
newArgs = append(newArgs, arg)
}
callExpr.Args = newArgs
return false
}
}
return true
})
return nil
}
// MainProviderSetAst 主 ProviderSet AST 注入结构体
// 用于向主层的 ProviderSet(如 internal/biz/biz.go)中注入新包的 Provider
type MainProviderSetAst struct {
Layer string // 层名称: biz, data, service
PackageName string // 新包名
ProviderName string // Provider 名称,如 NewArticleUsecase
ImportPath string // 导入路径
Path string // 目标文件路径
}
// NewMainProviderSetAst 创建主 ProviderSet AST 注入器
func NewMainProviderSetAst(layer, packageName, providerName, importPath, path string) *MainProviderSetAst {
return &MainProviderSetAst{
Layer: layer,
PackageName: packageName,
ProviderName: providerName,
ImportPath: importPath,
Path: path,
}
}
// Parse 解析文件
func (a *MainProviderSetAst) Parse(filename string, writer io.Writer) (*ast.File, error) {
if filename == "" {
filename = a.Path
}
if filename == "" {
return nil, fmt.Errorf("文件路径为空")
}
fset := token.NewFileSet()
file, err := parser.ParseFile(fset, filename, nil, parser.ParseComments)
if err != nil {
return nil, fmt.Errorf("解析文件失败 %s: %w", filename, err)
}
return file, nil
}
// Rollback 回滚注入
func (a *MainProviderSetAst) Rollback(file *ast.File) error {
// 移除导入
removeImport(file, a.ImportPath)
// 移除 Provider
return removeWireProviderWithPackage(file, a.PackageName, a.ProviderName)
}
// Injection 注入代码
func (a *MainProviderSetAst) Injection(file *ast.File) error {
// 添加导入
addImport(file, a.ImportPath)
// 添加 Provider
return addWireProviderWithPackage(file, a.PackageName, a.ProviderName)
}
// Format 格式化输出
func (a *MainProviderSetAst) Format(filename string, writer io.Writer, file *ast.File) error {
if filename == "" {
filename = a.Path
}
fset := token.NewFileSet()
if writer == nil {
f, err := os.OpenFile(filename, os.O_WRONLY|os.O_TRUNC, 0666)
if err != nil {
return errors.Wrapf(err, "打开文件失败: %s", filename)
}
defer f.Close()
writer = f
}
return format.Node(writer, fset, file)
}
// addImport 添加导入
func addImport(file *ast.File, importPath string) {
if importPath == "" {
return
}
impStr := fmt.Sprintf("\"%s\"", importPath)
// 检查是否已存在
for _, imp := range file.Imports {
if imp.Path.Value == impStr {
return
}
}
// 查找 import 声明并添加
for _, decl := range file.Decls {
genDecl, ok := decl.(*ast.GenDecl)
if !ok || genDecl.Tok != token.IMPORT {
continue
}
genDecl.Specs = append(genDecl.Specs, &ast.ImportSpec{
Path: &ast.BasicLit{
Kind: token.STRING,
Value: impStr,
},
})
return
}
// 如果没有 import 声明,创建一个新的
newDecl := &ast.GenDecl{
Tok: token.IMPORT,
Specs: []ast.Spec{
&ast.ImportSpec{
Path: &ast.BasicLit{
Kind: token.STRING,
Value: impStr,
},
},
},
}
file.Decls = append([]ast.Decl{newDecl}, file.Decls...)
}
// removeImport 移除导入
func removeImport(file *ast.File, importPath string) {
if importPath == "" {
return
}
impStr := fmt.Sprintf("\"%s\"", importPath)
for i := 0; i < len(file.Decls); i++ {
genDecl, ok := file.Decls[i].(*ast.GenDecl)
if !ok || genDecl.Tok != token.IMPORT {
continue
}
for j := 0; j < len(genDecl.Specs); j++ {
impSpec, ok := genDecl.Specs[j].(*ast.ImportSpec)
if !ok {
continue
}
if impSpec.Path.Value == impStr {
genDecl.Specs = append(genDecl.Specs[:j], genDecl.Specs[j+1:]...)
if len(genDecl.Specs) == 0 {
file.Decls = append(file.Decls[:i], file.Decls[i+1:]...)
}
return
}
}
}
}
// addWireProviderWithPackage 向 wire.NewSet 中添加带包名的 Provider
// 例如: article.NewArticleUsecase
func addWireProviderWithPackage(file *ast.File, packageName, providerName string) error {
var found bool
ast.Inspect(file, func(node ast.Node) bool {
genDecl, ok := node.(*ast.GenDecl)
if !ok || genDecl.Tok != token.VAR {
return true
}
for _, spec := range genDecl.Specs {
valueSpec, ok := spec.(*ast.ValueSpec)
if !ok {
continue
}
for i, name := range valueSpec.Names {
if name.Name != "ProviderSet" {
continue
}
if i >= len(valueSpec.Values) {
continue
}
callExpr, ok := valueSpec.Values[i].(*ast.CallExpr)
if !ok {
continue
}
selExpr, ok := callExpr.Fun.(*ast.SelectorExpr)
if !ok {
continue
}
ident, ok := selExpr.X.(*ast.Ident)
if !ok || ident.Name != "wire" || selExpr.Sel.Name != "NewSet" {
continue
}
// 构建完整的 Provider 名称
fullProviderName := fmt.Sprintf("%s.%s", packageName, providerName)
// 检查 Provider 是否已存在
for _, arg := range callExpr.Args {
if selArg, ok := arg.(*ast.SelectorExpr); ok {
if pkgIdent, ok := selArg.X.(*ast.Ident); ok {
existingName := fmt.Sprintf("%s.%s", pkgIdent.Name, selArg.Sel.Name)
if existingName == fullProviderName {
found = true
return false // 已存在,不需要添加
}
}
}
}
// 添加新的 Provider (package.ProviderName 格式)
newArg := &ast.SelectorExpr{
X: &ast.Ident{Name: packageName},
Sel: &ast.Ident{Name: providerName},
}
callExpr.Args = append(callExpr.Args, newArg)
found = true
return false
}
}
return true
})
if !found {
return fmt.Errorf("未找到 ProviderSet 变量或 wire.NewSet 调用")
}
return nil
}
// removeWireProviderWithPackage 从 wire.NewSet 中移除带包名的 Provider
func removeWireProviderWithPackage(file *ast.File, packageName, providerName string) error {
fullProviderName := fmt.Sprintf("%s.%s", packageName, providerName)
ast.Inspect(file, func(node ast.Node) bool {
genDecl, ok := node.(*ast.GenDecl)
if !ok || genDecl.Tok != token.VAR {
return true
}
for _, spec := range genDecl.Specs {
valueSpec, ok := spec.(*ast.ValueSpec)
if !ok {
continue
}
for i, name := range valueSpec.Names {
if name.Name != "ProviderSet" {
continue
}
if i >= len(valueSpec.Values) {
continue
}
callExpr, ok := valueSpec.Values[i].(*ast.CallExpr)
if !ok {
continue
}
selExpr, ok := callExpr.Fun.(*ast.SelectorExpr)
if !ok {
continue
}
ident, ok := selExpr.X.(*ast.Ident)
if !ok || ident.Name != "wire" || selExpr.Sel.Name != "NewSet" {
continue
}
// 移除指定的 Provider
newArgs := make([]ast.Expr, 0, len(callExpr.Args))
for _, arg := range callExpr.Args {
if selArg, ok := arg.(*ast.SelectorExpr); ok {
if pkgIdent, ok := selArg.X.(*ast.Ident); ok {
existingName := fmt.Sprintf("%s.%s", pkgIdent.Name, selArg.Sel.Name)
if existingName == fullProviderName {
continue // 跳过要移除的 Provider
}
}
}
newArgs = append(newArgs, arg)
}
callExpr.Args = newArgs
return false
}
}
return true
})
return nil
}
// GetLayerMainFile 获取层的主文件路径
func GetLayerMainFile(layer string) string {
switch layer {
case "biz":
return "internal/biz/biz.go"
case "data":
return "internal/data/data.go"
case "service":
return "internal/service/service.go"
case "handler", "api":
return "internal/server/handler/enter.go"
case "router":
return "internal/server/router/enter.go"
default:
return ""
}
}
// GetLayerImportPath 获取层的导入路径
func GetLayerImportPath(module, layer, packageName string) string {
switch layer {
case "biz":
return fmt.Sprintf("%s/internal/biz/%s", module, packageName)
case "data":
return fmt.Sprintf("%s/internal/data/%s", module, packageName)
case "service":
return fmt.Sprintf("%s/internal/service/%s", module, packageName)
case "handler", "api":
return fmt.Sprintf("%s/internal/server/handler/%s", module, packageName)
case "router":
return fmt.Sprintf("%s/internal/server/router/%s", module, packageName)
default:
return ""
}
}
// GetProviderName 获取 Provider 名称
func GetProviderName(layer, structName string) string {
switch layer {
case "biz":
return fmt.Sprintf("New%sUsecase", structName)
case "data":
return fmt.Sprintf("New%sRepo", structName)
case "service":
return fmt.Sprintf("New%sService", structName)
case "handler", "api":
return fmt.Sprintf("New%sHandler", structName)
case "router":
return fmt.Sprintf("New%sRouter", structName)
default:
return ""
}
}
// InjectWireProvider 注入 Wire Provider 到主文件
// 这是一个便捷函数,用于一次性完成导入和 Provider 注入
func InjectWireProvider(module, layer, packageName, structName string) error {
mainFile := GetLayerMainFile(layer)
if mainFile == "" {
return fmt.Errorf("未知的层: %s", layer)
}
importPath := GetLayerImportPath(module, layer, packageName)
providerName := GetProviderName(layer, structName)
injector := NewMainProviderSetAst(layer, packageName, providerName, importPath, mainFile)
file, err := injector.Parse("", nil)
if err != nil {
return err
}
if err := injector.Injection(file); err != nil {
return err
}
return injector.Format("", nil, file)
}
// RollbackWireProvider 回滚 Wire Provider 注入
func RollbackWireProvider(module, layer, packageName, structName string) error {
mainFile := GetLayerMainFile(layer)
if mainFile == "" {
return fmt.Errorf("未知的层: %s", layer)
}
importPath := GetLayerImportPath(module, layer, packageName)
providerName := GetProviderName(layer, structName)
injector := NewMainProviderSetAst(layer, packageName, providerName, importPath, mainFile)
file, err := injector.Parse("", nil)
if err != nil {
return err
}
if err := injector.Rollback(file); err != nil {
return err
}
return injector.Format("", nil, file)
}
// WireInjectionConfig Wire 注入配置
type WireInjectionConfig struct {
Module string // 模块名 (go.mod 中的 module)
PackageName string // 包名
StructName string // 结构体名
Layers []string // 需要注入的层
}
// InjectAllWireProviders 注入所有层的 Wire Provider
func InjectAllWireProviders(config *WireInjectionConfig) error {
for _, layer := range config.Layers {
if err := InjectWireProvider(config.Module, layer, config.PackageName, config.StructName); err != nil {
// 如果注入失败,尝试回滚已注入的
for _, rollbackLayer := range config.Layers {
if rollbackLayer == layer {
break
}
_ = RollbackWireProvider(config.Module, rollbackLayer, config.PackageName, config.StructName)
}
return fmt.Errorf("注入 %s 层 Wire Provider 失败: %w", layer, err)
}
}
return nil
}
// RollbackAllWireProviders 回滚所有层的 Wire Provider
func RollbackAllWireProviders(config *WireInjectionConfig) error {
var lastErr error
for _, layer := range config.Layers {
if err := RollbackWireProvider(config.Module, layer, config.PackageName, config.StructName); err != nil {
lastErr = err
}
}
return lastErr
}
+684
View File
@@ -0,0 +1,684 @@
package autocode
import (
"bytes"
"go/ast"
"go/parser"
"go/token"
"strings"
"testing"
)
// TestWireProviderAst_Parse 测试 WireProviderAst 解析功能
func TestWireProviderAst_Parse(t *testing.T) {
// 创建测试用的 Go 代码
testCode := `package biz
import "github.com/google/wire"
var ProviderSet = wire.NewSet(NewUserUsecase)
`
// 创建临时文件
fset := token.NewFileSet()
file, err := parser.ParseFile(fset, "test.go", testCode, parser.ParseComments)
if err != nil {
t.Fatalf("解析测试代码失败: %v", err)
}
if file == nil {
t.Fatal("解析结果为空")
}
if file.Name.Name != "biz" {
t.Errorf("包名应为 'biz', 实际为 '%s'", file.Name.Name)
}
}
// TestWireProviderAst_Injection 测试 WireProviderAst 注入功能
func TestWireProviderAst_Injection(t *testing.T) {
testCode := `package biz
import "github.com/google/wire"
var ProviderSet = wire.NewSet(NewUserUsecase)
`
fset := token.NewFileSet()
file, err := parser.ParseFile(fset, "test.go", testCode, parser.ParseComments)
if err != nil {
t.Fatalf("解析测试代码失败: %v", err)
}
// 创建注入器
injector := &WireProviderAst{
Layer: "biz",
PackageName: "article",
ProviderName: "NewArticleUsecase",
Path: "test.go",
}
// 执行注入
err = injector.Injection(file)
if err != nil {
t.Fatalf("注入失败: %v", err)
}
// 验证注入结果
var found bool
ast.Inspect(file, func(node ast.Node) bool {
if genDecl, ok := node.(*ast.GenDecl); ok && genDecl.Tok == token.VAR {
for _, spec := range genDecl.Specs {
if valueSpec, ok := spec.(*ast.ValueSpec); ok {
for i, name := range valueSpec.Names {
if name.Name == "ProviderSet" && i < len(valueSpec.Values) {
if callExpr, ok := valueSpec.Values[i].(*ast.CallExpr); ok {
for _, arg := range callExpr.Args {
if ident, ok := arg.(*ast.Ident); ok && ident.Name == "NewArticleUsecase" {
found = true
return false
}
}
}
}
}
}
}
}
return true
})
if !found {
t.Error("注入后应包含 NewArticleUsecase Provider")
}
}
// TestWireProviderAst_Injection_Duplicate 测试重复注入不会添加重复项
func TestWireProviderAst_Injection_Duplicate(t *testing.T) {
testCode := `package biz
import "github.com/google/wire"
var ProviderSet = wire.NewSet(NewUserUsecase, NewArticleUsecase)
`
fset := token.NewFileSet()
file, err := parser.ParseFile(fset, "test.go", testCode, parser.ParseComments)
if err != nil {
t.Fatalf("解析测试代码失败: %v", err)
}
injector := &WireProviderAst{
Layer: "biz",
PackageName: "article",
ProviderName: "NewArticleUsecase",
Path: "test.go",
}
// 执行注入(应该不会添加重复项)
err = injector.Injection(file)
if err != nil {
t.Fatalf("注入失败: %v", err)
}
// 统计 NewArticleUsecase 出现次数
count := 0
ast.Inspect(file, func(node ast.Node) bool {
if genDecl, ok := node.(*ast.GenDecl); ok && genDecl.Tok == token.VAR {
for _, spec := range genDecl.Specs {
if valueSpec, ok := spec.(*ast.ValueSpec); ok {
for i, name := range valueSpec.Names {
if name.Name == "ProviderSet" && i < len(valueSpec.Values) {
if callExpr, ok := valueSpec.Values[i].(*ast.CallExpr); ok {
for _, arg := range callExpr.Args {
if ident, ok := arg.(*ast.Ident); ok && ident.Name == "NewArticleUsecase" {
count++
}
}
}
}
}
}
}
}
return true
})
if count != 1 {
t.Errorf("NewArticleUsecase 应该只出现 1 次, 实际出现 %d 次", count)
}
}
// TestWireProviderAst_Rollback 测试 WireProviderAst 回滚功能
func TestWireProviderAst_Rollback(t *testing.T) {
testCode := `package biz
import "github.com/google/wire"
var ProviderSet = wire.NewSet(NewUserUsecase, NewArticleUsecase)
`
fset := token.NewFileSet()
file, err := parser.ParseFile(fset, "test.go", testCode, parser.ParseComments)
if err != nil {
t.Fatalf("解析测试代码失败: %v", err)
}
injector := &WireProviderAst{
Layer: "biz",
PackageName: "article",
ProviderName: "NewArticleUsecase",
Path: "test.go",
}
// 执行回滚
err = injector.Rollback(file)
if err != nil {
t.Fatalf("回滚失败: %v", err)
}
// 验证回滚结果
var found bool
ast.Inspect(file, func(node ast.Node) bool {
if genDecl, ok := node.(*ast.GenDecl); ok && genDecl.Tok == token.VAR {
for _, spec := range genDecl.Specs {
if valueSpec, ok := spec.(*ast.ValueSpec); ok {
for i, name := range valueSpec.Names {
if name.Name == "ProviderSet" && i < len(valueSpec.Values) {
if callExpr, ok := valueSpec.Values[i].(*ast.CallExpr); ok {
for _, arg := range callExpr.Args {
if ident, ok := arg.(*ast.Ident); ok && ident.Name == "NewArticleUsecase" {
found = true
return false
}
}
}
}
}
}
}
}
return true
})
if found {
t.Error("回滚后不应包含 NewArticleUsecase Provider")
}
}
// TestWireProviderAst_Format 测试 WireProviderAst 格式化输出
func TestWireProviderAst_Format(t *testing.T) {
testCode := `package biz
import "github.com/google/wire"
var ProviderSet = wire.NewSet(NewUserUsecase)
`
fset := token.NewFileSet()
file, err := parser.ParseFile(fset, "test.go", testCode, parser.ParseComments)
if err != nil {
t.Fatalf("解析测试代码失败: %v", err)
}
injector := &WireProviderAst{
Layer: "biz",
PackageName: "article",
ProviderName: "NewArticleUsecase",
Path: "test.go",
}
// 注入后格式化
_ = injector.Injection(file)
var buf bytes.Buffer
err = injector.Format("", &buf, file)
if err != nil {
t.Fatalf("格式化失败: %v", err)
}
result := buf.String()
if !strings.Contains(result, "NewArticleUsecase") {
t.Error("格式化输出应包含 NewArticleUsecase")
}
if !strings.Contains(result, "package biz") {
t.Error("格式化输出应包含 package biz")
}
}
// TestMainProviderSetAst_Injection 测试 MainProviderSetAst 注入功能
func TestMainProviderSetAst_Injection(t *testing.T) {
testCode := `package biz
import (
"github.com/google/wire"
"kra/internal/biz/user"
)
var ProviderSet = wire.NewSet(user.NewUserUsecase)
`
fset := token.NewFileSet()
file, err := parser.ParseFile(fset, "test.go", testCode, parser.ParseComments)
if err != nil {
t.Fatalf("解析测试代码失败: %v", err)
}
injector := &MainProviderSetAst{
Layer: "biz",
PackageName: "article",
ProviderName: "NewArticleUsecase",
ImportPath: "kra/internal/biz/article",
Path: "test.go",
}
// 执行注入
err = injector.Injection(file)
if err != nil {
t.Fatalf("注入失败: %v", err)
}
// 验证导入是否添加 - 检查 import 声明中是否包含新导入
var importFound bool
ast.Inspect(file, func(node ast.Node) bool {
if genDecl, ok := node.(*ast.GenDecl); ok && genDecl.Tok == token.IMPORT {
for _, spec := range genDecl.Specs {
if impSpec, ok := spec.(*ast.ImportSpec); ok {
if impSpec.Path.Value == "\"kra/internal/biz/article\"" {
importFound = true
return false
}
}
}
}
return true
})
if !importFound {
t.Error("注入后应包含 kra/internal/biz/article 导入")
}
// 验证 Provider 是否添加
var providerFound bool
ast.Inspect(file, func(node ast.Node) bool {
if genDecl, ok := node.(*ast.GenDecl); ok && genDecl.Tok == token.VAR {
for _, spec := range genDecl.Specs {
if valueSpec, ok := spec.(*ast.ValueSpec); ok {
for i, name := range valueSpec.Names {
if name.Name == "ProviderSet" && i < len(valueSpec.Values) {
if callExpr, ok := valueSpec.Values[i].(*ast.CallExpr); ok {
for _, arg := range callExpr.Args {
if selExpr, ok := arg.(*ast.SelectorExpr); ok {
if pkgIdent, ok := selExpr.X.(*ast.Ident); ok {
if pkgIdent.Name == "article" && selExpr.Sel.Name == "NewArticleUsecase" {
providerFound = true
return false
}
}
}
}
}
}
}
}
}
}
return true
})
if !providerFound {
t.Error("注入后应包含 article.NewArticleUsecase Provider")
}
}
// TestMainProviderSetAst_Rollback 测试 MainProviderSetAst 回滚功能
func TestMainProviderSetAst_Rollback(t *testing.T) {
testCode := `package biz
import (
"github.com/google/wire"
"kra/internal/biz/user"
"kra/internal/biz/article"
)
var ProviderSet = wire.NewSet(user.NewUserUsecase, article.NewArticleUsecase)
`
fset := token.NewFileSet()
file, err := parser.ParseFile(fset, "test.go", testCode, parser.ParseComments)
if err != nil {
t.Fatalf("解析测试代码失败: %v", err)
}
injector := &MainProviderSetAst{
Layer: "biz",
PackageName: "article",
ProviderName: "NewArticleUsecase",
ImportPath: "kra/internal/biz/article",
Path: "test.go",
}
// 执行回滚
err = injector.Rollback(file)
if err != nil {
t.Fatalf("回滚失败: %v", err)
}
// 验证导入是否移除 - 检查 import 声明
var importFound bool
ast.Inspect(file, func(node ast.Node) bool {
if genDecl, ok := node.(*ast.GenDecl); ok && genDecl.Tok == token.IMPORT {
for _, spec := range genDecl.Specs {
if impSpec, ok := spec.(*ast.ImportSpec); ok {
if impSpec.Path.Value == "\"kra/internal/biz/article\"" {
importFound = true
return false
}
}
}
}
return true
})
if importFound {
t.Error("回滚后不应包含 kra/internal/biz/article 导入")
}
// 验证 Provider 是否移除
var providerFound bool
ast.Inspect(file, func(node ast.Node) bool {
if genDecl, ok := node.(*ast.GenDecl); ok && genDecl.Tok == token.VAR {
for _, spec := range genDecl.Specs {
if valueSpec, ok := spec.(*ast.ValueSpec); ok {
for i, name := range valueSpec.Names {
if name.Name == "ProviderSet" && i < len(valueSpec.Values) {
if callExpr, ok := valueSpec.Values[i].(*ast.CallExpr); ok {
for _, arg := range callExpr.Args {
if selExpr, ok := arg.(*ast.SelectorExpr); ok {
if pkgIdent, ok := selExpr.X.(*ast.Ident); ok {
if pkgIdent.Name == "article" && selExpr.Sel.Name == "NewArticleUsecase" {
providerFound = true
return false
}
}
}
}
}
}
}
}
}
}
return true
})
if providerFound {
t.Error("回滚后不应包含 article.NewArticleUsecase Provider")
}
}
// TestGetLayerMainFile 测试获取层主文件路径
func TestGetLayerMainFile(t *testing.T) {
testCases := []struct {
layer string
expected string
}{
{"biz", "internal/biz/biz.go"},
{"data", "internal/data/data.go"},
{"service", "internal/service/service.go"},
{"handler", "internal/server/handler/enter.go"},
{"api", "internal/server/handler/enter.go"},
{"router", "internal/server/router/enter.go"},
{"unknown", ""},
}
for _, tc := range testCases {
t.Run(tc.layer, func(t *testing.T) {
result := GetLayerMainFile(tc.layer)
if result != tc.expected {
t.Errorf("GetLayerMainFile(%s) = %s, 期望 %s", tc.layer, result, tc.expected)
}
})
}
}
// TestGetLayerImportPath 测试获取层导入路径
func TestGetLayerImportPath(t *testing.T) {
module := "kra"
packageName := "article"
testCases := []struct {
layer string
expected string
}{
{"biz", "kra/internal/biz/article"},
{"data", "kra/internal/data/article"},
{"service", "kra/internal/service/article"},
{"handler", "kra/internal/server/handler/article"},
{"api", "kra/internal/server/handler/article"},
{"router", "kra/internal/server/router/article"},
{"unknown", ""},
}
for _, tc := range testCases {
t.Run(tc.layer, func(t *testing.T) {
result := GetLayerImportPath(module, tc.layer, packageName)
if result != tc.expected {
t.Errorf("GetLayerImportPath(%s, %s, %s) = %s, 期望 %s", module, tc.layer, packageName, result, tc.expected)
}
})
}
}
// TestGetProviderName 测试获取 Provider 名称
func TestGetProviderName(t *testing.T) {
structName := "Article"
testCases := []struct {
layer string
expected string
}{
{"biz", "NewArticleUsecase"},
{"data", "NewArticleRepo"},
{"service", "NewArticleService"},
{"handler", "NewArticleHandler"},
{"api", "NewArticleHandler"},
{"router", "NewArticleRouter"},
{"unknown", ""},
}
for _, tc := range testCases {
t.Run(tc.layer, func(t *testing.T) {
result := GetProviderName(tc.layer, structName)
if result != tc.expected {
t.Errorf("GetProviderName(%s, %s) = %s, 期望 %s", tc.layer, structName, result, tc.expected)
}
})
}
}
// TestAddImport 测试添加导入
func TestAddImport(t *testing.T) {
testCode := `package main
import "fmt"
func main() {}
`
fset := token.NewFileSet()
file, err := parser.ParseFile(fset, "test.go", testCode, parser.ParseComments)
if err != nil {
t.Fatalf("解析测试代码失败: %v", err)
}
// 添加新导入
addImport(file, "kra/internal/biz/article")
// 验证导入是否添加 - 检查 import 声明
var found bool
ast.Inspect(file, func(node ast.Node) bool {
if genDecl, ok := node.(*ast.GenDecl); ok && genDecl.Tok == token.IMPORT {
for _, spec := range genDecl.Specs {
if impSpec, ok := spec.(*ast.ImportSpec); ok {
if impSpec.Path.Value == "\"kra/internal/biz/article\"" {
found = true
return false
}
}
}
}
return true
})
if !found {
t.Error("应该添加 kra/internal/biz/article 导入")
}
}
// TestAddImport_Duplicate 测试重复添加导入
func TestAddImport_Duplicate(t *testing.T) {
testCode := `package main
import (
"fmt"
"kra/internal/biz/article"
)
func main() {}
`
fset := token.NewFileSet()
file, err := parser.ParseFile(fset, "test.go", testCode, parser.ParseComments)
if err != nil {
t.Fatalf("解析测试代码失败: %v", err)
}
initialCount := len(file.Imports)
// 尝试添加已存在的导入
addImport(file, "kra/internal/biz/article")
// 验证导入数量没有增加
if len(file.Imports) != initialCount {
t.Errorf("重复添加导入后数量应保持 %d, 实际为 %d", initialCount, len(file.Imports))
}
}
// TestRemoveImport 测试移除导入
func TestRemoveImport(t *testing.T) {
testCode := `package main
import (
"fmt"
"kra/internal/biz/article"
)
func main() {}
`
fset := token.NewFileSet()
file, err := parser.ParseFile(fset, "test.go", testCode, parser.ParseComments)
if err != nil {
t.Fatalf("解析测试代码失败: %v", err)
}
// 移除导入
removeImport(file, "kra/internal/biz/article")
// 验证导入是否移除 - 检查 import 声明
var found bool
ast.Inspect(file, func(node ast.Node) bool {
if genDecl, ok := node.(*ast.GenDecl); ok && genDecl.Tok == token.IMPORT {
for _, spec := range genDecl.Specs {
if impSpec, ok := spec.(*ast.ImportSpec); ok {
if impSpec.Path.Value == "\"kra/internal/biz/article\"" {
found = true
return false
}
}
}
}
return true
})
if found {
t.Error("应该移除 kra/internal/biz/article 导入")
}
}
// TestNewWireProviderAst 测试创建 WireProviderAst
func TestNewWireProviderAst(t *testing.T) {
ast := NewWireProviderAst("biz", "article", "NewArticleUsecase", "internal/biz/article/enter.go")
if ast.Layer != "biz" {
t.Errorf("Layer 应为 'biz', 实际为 '%s'", ast.Layer)
}
if ast.PackageName != "article" {
t.Errorf("PackageName 应为 'article', 实际为 '%s'", ast.PackageName)
}
if ast.ProviderName != "NewArticleUsecase" {
t.Errorf("ProviderName 应为 'NewArticleUsecase', 实际为 '%s'", ast.ProviderName)
}
if ast.Path != "internal/biz/article/enter.go" {
t.Errorf("Path 应为 'internal/biz/article/enter.go', 实际为 '%s'", ast.Path)
}
}
// TestNewMainProviderSetAst 测试创建 MainProviderSetAst
func TestNewMainProviderSetAst(t *testing.T) {
ast := NewMainProviderSetAst("biz", "article", "NewArticleUsecase", "kra/internal/biz/article", "internal/biz/biz.go")
if ast.Layer != "biz" {
t.Errorf("Layer 应为 'biz', 实际为 '%s'", ast.Layer)
}
if ast.PackageName != "article" {
t.Errorf("PackageName 应为 'article', 实际为 '%s'", ast.PackageName)
}
if ast.ProviderName != "NewArticleUsecase" {
t.Errorf("ProviderName 应为 'NewArticleUsecase', 实际为 '%s'", ast.ProviderName)
}
if ast.ImportPath != "kra/internal/biz/article" {
t.Errorf("ImportPath 应为 'kra/internal/biz/article', 实际为 '%s'", ast.ImportPath)
}
if ast.Path != "internal/biz/biz.go" {
t.Errorf("Path 应为 'internal/biz/biz.go', 实际为 '%s'", ast.Path)
}
}
// TestWireProviderAst_Injection_NoProviderSet 测试没有 ProviderSet 时的注入
func TestWireProviderAst_Injection_NoProviderSet(t *testing.T) {
testCode := `package biz
import "github.com/google/wire"
var SomeOtherVar = "test"
`
fset := token.NewFileSet()
file, err := parser.ParseFile(fset, "test.go", testCode, parser.ParseComments)
if err != nil {
t.Fatalf("解析测试代码失败: %v", err)
}
injector := &WireProviderAst{
Layer: "biz",
PackageName: "article",
ProviderName: "NewArticleUsecase",
Path: "test.go",
}
// 执行注入应该返回错误
err = injector.Injection(file)
if err == nil {
t.Error("没有 ProviderSet 时注入应该返回错误")
}
}
// TestWireInjectionConfig 测试 WireInjectionConfig 结构体
func TestWireInjectionConfig(t *testing.T) {
config := &WireInjectionConfig{
Module: "kra",
PackageName: "article",
StructName: "Article",
Layers: []string{"biz", "data", "service", "handler", "router"},
}
if config.Module != "kra" {
t.Errorf("Module 应为 'kra', 实际为 '%s'", config.Module)
}
if config.PackageName != "article" {
t.Errorf("PackageName 应为 'article', 实际为 '%s'", config.PackageName)
}
if config.StructName != "Article" {
t.Errorf("StructName 应为 'Article', 实际为 '%s'", config.StructName)
}
if len(config.Layers) != 5 {
t.Errorf("Layers 长度应为 5, 实际为 %d", len(config.Layers))
}
}
File diff suppressed because it is too large Load Diff
+34
View File
@@ -1,6 +1,7 @@
package utils
import (
"context"
"net"
"time"
@@ -8,6 +9,7 @@ import (
"github.com/gin-gonic/gin"
"github.com/google/uuid"
"github.com/redis/go-redis/v9"
"go.uber.org/zap"
)
@@ -17,6 +19,12 @@ var jwtInstance *jwt.JWT
// 全局日志实例(需要在初始化时设置)
var logger *zap.Logger
// 全局 Redis 客户端(需要在初始化时设置)
var redisClient redis.UniversalClient
// JWT 过期时间
var jwtExpiresTime time.Duration = 7 * 24 * time.Hour
// SetJWT 设置全局 JWT 实例
func SetJWT(j *jwt.JWT) {
jwtInstance = j
@@ -27,6 +35,32 @@ func SetLogger(l *zap.Logger) {
logger = l
}
// SetRedisClient 设置全局 Redis 客户端
func SetRedisClient(client redis.UniversalClient) {
redisClient = client
}
// SetJWTExpiresTime 设置 JWT 过期时间
func SetJWTExpiresTime(d time.Duration) {
jwtExpiresTime = d
}
// SetRedisJWT 设置 Redis JWT
func SetRedisJWT(token, username string) error {
if redisClient == nil {
return nil
}
return redisClient.Set(context.Background(), "jwt:"+username, token, jwtExpiresTime).Err()
}
// GetRedisJWT 从 Redis 获取 JWT
func GetRedisJWT(username string) (string, error) {
if redisClient == nil {
return "", redis.Nil
}
return redisClient.Get(context.Background(), "jwt:"+username).Result()
}
// ClearToken 清除 token
func ClearToken(c *gin.Context) {
// 增加cookie x-token 向来源的web添加
+74
View File
@@ -0,0 +1,74 @@
package utils
import (
"io"
"os"
"path/filepath"
)
// CopyFile 复制文件
func CopyFile(src, dst string) error {
srcFile, err := os.Open(src)
if err != nil {
return err
}
defer srcFile.Close()
// 确保目标目录存在
dstDir := filepath.Dir(dst)
if err := os.MkdirAll(dstDir, os.ModePerm); err != nil {
return err
}
dstFile, err := os.Create(dst)
if err != nil {
return err
}
defer dstFile.Close()
_, err = io.Copy(dstFile, srcFile)
return err
}
// CopyDir 复制目录
func CopyDir(src, dst string) error {
srcInfo, err := os.Stat(src)
if err != nil {
return err
}
if err = os.MkdirAll(dst, srcInfo.Mode()); err != nil {
return err
}
entries, err := os.ReadDir(src)
if err != nil {
return err
}
for _, entry := range entries {
srcPath := filepath.Join(src, entry.Name())
dstPath := filepath.Join(dst, entry.Name())
if entry.IsDir() {
if err = CopyDir(srcPath, dstPath); err != nil {
return err
}
} else {
if err = CopyFile(srcPath, dstPath); err != nil {
return err
}
}
}
return nil
}
// RemoveFile 删除文件
func RemoveFile(path string) error {
return os.Remove(path)
}
// RemoveDir 删除目录
func RemoveDir(path string) error {
return os.RemoveAll(path)
}
+125
View File
@@ -0,0 +1,125 @@
package utils
import (
"fmt"
"math/rand"
"reflect"
"strings"
)
// StructToMap 利用反射将结构体转化为map
func StructToMap(obj interface{}) map[string]interface{} {
obj1 := reflect.TypeOf(obj)
obj2 := reflect.ValueOf(obj)
data := make(map[string]interface{})
for i := 0; i < obj1.NumField(); i++ {
if obj1.Field(i).Tag.Get("mapstructure") != "" {
data[obj1.Field(i).Tag.Get("mapstructure")] = obj2.Field(i).Interface()
} else {
data[obj1.Field(i).Name] = obj2.Field(i).Interface()
}
}
return data
}
// ArrayToString 将数组格式化为字符串
func ArrayToString(array []interface{}) string {
return strings.Replace(strings.Trim(fmt.Sprint(array), "[]"), " ", ",", -1)
}
// Pointer 返回值的指针
func Pointer[T any](in T) (out *T) {
return &in
}
// FirstUpper 首字母大写
func FirstUpper(s string) string {
if s == "" {
return ""
}
return strings.ToUpper(s[:1]) + s[1:]
}
// FirstLower 首字母小写
func FirstLower(s string) string {
if s == "" {
return ""
}
return strings.ToLower(s[:1]) + s[1:]
}
// MaheHump 将字符串转换为驼峰命名
func MaheHump(s string) string {
words := strings.Split(s, "-")
for i := 1; i < len(words); i++ {
words[i] = strings.Title(words[i])
}
return strings.Join(words, "")
}
// HumpToUnderscore 将驼峰命名转换为下划线分割模式
func HumpToUnderscore(s string) string {
var result strings.Builder
for i, char := range s {
if i > 0 && char >= 'A' && char <= 'Z' {
// 在大写字母前添加下划线
result.WriteRune('_')
result.WriteRune(char - 'A' + 'a') // 转小写
} else {
result.WriteRune(char)
}
}
return strings.ToLower(result.String())
}
// RandomString 随机字符串
func RandomString(n int) string {
var letters = []rune("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789")
b := make([]rune, n)
for i := range b {
b[i] = letters[RandomInt(0, len(letters))]
}
return string(b)
}
// RandomInt 生成指定范围内的随机整数
func RandomInt(min, max int) int {
return min + rand.Intn(max-min)
}
// TreeNode 树节点接口
type TreeNode[T any] interface {
GetID() int
GetParentID() int
SetChildren(child T)
}
// BuildTree 用于构建一个树形结构
func BuildTree[T TreeNode[T]](nodes []T) []T {
nodeMap := make(map[int]T)
// 创建一个基本map
for i := range nodes {
nodeMap[nodes[i].GetID()] = nodes[i]
}
for i := range nodes {
if nodes[i].GetParentID() != 0 {
parent := nodeMap[nodes[i].GetParentID()]
parent.SetChildren(nodes[i])
}
}
var rootNodes []T
for i := range nodeMap {
if nodeMap[i].GetParentID() == 0 {
rootNodes = append(rootNodes, nodeMap[i])
}
}
return rootNodes
}
+17
View File
@@ -0,0 +1,17 @@
package utils
import (
"golang.org/x/crypto/bcrypt"
)
// BcryptHash 使用 bcrypt 对密码进行加密
func BcryptHash(password string) string {
bytes, _ := bcrypt.GenerateFromPassword([]byte(password), bcrypt.DefaultCost)
return string(bytes)
}
// BcryptCheck 对比明文密码和数据库的哈希值
func BcryptCheck(password, hash string) bool {
err := bcrypt.CompareHashAndPassword([]byte(hash), []byte(password))
return err == nil
}
+30
View File
@@ -0,0 +1,30 @@
package utils
import (
"strconv"
"strings"
"time"
)
// ParseDuration 解析时间字符串,支持天数格式(如 "7d" 或 "7d12h")
func ParseDuration(d string) (time.Duration, error) {
d = strings.TrimSpace(d)
dr, err := time.ParseDuration(d)
if err == nil {
return dr, nil
}
if strings.Contains(d, "d") {
index := strings.Index(d, "d")
hour, _ := strconv.Atoi(d[:index])
dr = time.Hour * 24 * time.Duration(hour)
ndr, err := time.ParseDuration(d[index+1:])
if err != nil {
return dr, nil
}
return dr + ndr, nil
}
dv, err := strconv.ParseInt(d, 10, 64)
return time.Duration(dv), err
}
+3 -3
View File
@@ -6,10 +6,10 @@ import (
)
// MD5V 计算字节数组的MD5值
func MD5V(data []byte) string {
func MD5V(str []byte, b ...byte) string {
h := md5.New()
h.Write(data)
return hex.EncodeToString(h.Sum(nil))
h.Write(str)
return hex.EncodeToString(h.Sum(b))
}
// MD5String 计算字符串的MD5值
+18
View File
@@ -0,0 +1,18 @@
package plugin
import (
"github.com/gin-gonic/gin"
)
const (
OnlyFuncName = "Plugin"
)
// Plugin 插件模式接口化
type Plugin interface {
// Register 注册路由
Register(group *gin.RouterGroup)
// RouterPath 用户返回注册路由
RouterPath() string
}
+66
View File
@@ -0,0 +1,66 @@
package tool
import (
"github.com/pkg/errors"
"go.uber.org/zap"
"gorm.io/gorm"
)
// APIEntity API实体接口
type APIEntity interface {
GetPath() string
GetMethod() string
GetAPIGroup() string
}
// MenuEntity 菜单实体接口
type MenuEntity interface {
GetName() string
SetParentID(id int64)
}
// RegisterApis 注册API到数据库(通用版本)
func RegisterApis[T APIEntity](db *gorm.DB, model T, apis ...T) {
err := db.Transaction(func(tx *gorm.DB) error {
for _, api := range apis {
err := tx.Model(model).Where("path = ? AND method = ? AND api_group = ? ", api.GetPath(), api.GetMethod(), api.GetAPIGroup()).FirstOrCreate(&api).Error
if err != nil {
zap.L().Error("注册API失败", zap.Error(err), zap.String("api", api.GetPath()), zap.String("method", api.GetMethod()), zap.String("apiGroup", api.GetAPIGroup()))
return err
}
}
return nil
})
if err != nil {
zap.L().Error("注册API失败", zap.Error(err))
}
}
// RegisterMenus 注册菜单到数据库(通用版本)
func RegisterMenus[T MenuEntity](db *gorm.DB, model T, menus ...T) {
if len(menus) == 0 {
return
}
parentMenu := menus[0]
otherMenus := menus[1:]
err := db.Transaction(func(tx *gorm.DB) error {
err := tx.Model(model).Where("name = ? ", parentMenu.GetName()).FirstOrCreate(&parentMenu).Error
if err != nil {
zap.L().Error("注册菜单失败", zap.Error(err))
return errors.Wrap(err, "注册菜单失败")
}
// 获取父菜单ID需要通过反射或其他方式
for i := range otherMenus {
err = tx.Model(model).Where("name = ? ", otherMenus[i].GetName()).FirstOrCreate(&otherMenus[i]).Error
if err != nil {
zap.L().Error("注册菜单失败", zap.Error(err))
return errors.Wrap(err, "注册菜单失败")
}
}
return nil
})
if err != nil {
zap.L().Error("注册菜单失败", zap.Error(err))
}
}
+11
View File
@@ -0,0 +1,11 @@
package plugin
import (
"github.com/gin-gonic/gin"
)
// Plugin 插件模式接口化v2
type Plugin interface {
// Register 注册路由
Register(group *gin.Engine)
}
-19
View File
@@ -1,19 +0,0 @@
package utils
import (
"math/rand"
"time"
)
const charset = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"
var seededRand = rand.New(rand.NewSource(time.Now().UnixNano()))
// RandomString 生成指定长度的随机字符串
func RandomString(length int) string {
b := make([]byte, length)
for i := range b {
b[i] = charset[seededRand.Intn(len(charset))]
}
return string(b)
}
+68
View File
@@ -0,0 +1,68 @@
package request
import (
"bytes"
"encoding/json"
"net/http"
"net/url"
)
// HttpRequest 发送HTTP请求的通用方法
// urlStr: 请求URL
// method: HTTP方法 (GET, POST, PUT, DELETE等)
// headers: 请求头
// params: URL查询参数
// data: 请求体数据 (会被JSON序列化)
func HttpRequest(
urlStr string,
method string,
headers map[string]string,
params map[string]string,
data any) (*http.Response, error) {
// 创建URL
u, err := url.Parse(urlStr)
if err != nil {
return nil, err
}
// 添加查询参数
query := u.Query()
for k, v := range params {
query.Set(k, v)
}
u.RawQuery = query.Encode()
// 将数据编码为JSON
buf := new(bytes.Buffer)
if data != nil {
b, err := json.Marshal(data)
if err != nil {
return nil, err
}
buf = bytes.NewBuffer(b)
}
// 创建请求
req, err := http.NewRequest(method, u.String(), buf)
if err != nil {
return nil, err
}
for k, v := range headers {
req.Header.Set(k, v)
}
if data != nil {
req.Header.Set("Content-Type", "application/json")
}
// 发送请求
resp, err := http.DefaultClient.Do(req)
if err != nil {
return nil, err
}
// 返回响应,让调用者处理
return resp, nil
}
+108
View File
@@ -0,0 +1,108 @@
package utils
import (
"runtime"
"time"
"github.com/shirou/gopsutil/v3/cpu"
"github.com/shirou/gopsutil/v3/disk"
"github.com/shirou/gopsutil/v3/mem"
)
const (
B = 1
KB = 1024 * B
MB = 1024 * KB
GB = 1024 * MB
)
type Server struct {
Os Os `json:"os"`
Cpu Cpu `json:"cpu"`
Ram Ram `json:"ram"`
Disk []Disk `json:"disk"`
}
type Os struct {
GOOS string `json:"goos"`
NumCPU int `json:"numCpu"`
Compiler string `json:"compiler"`
GoVersion string `json:"goVersion"`
NumGoroutine int `json:"numGoroutine"`
}
type Cpu struct {
Cpus []float64 `json:"cpus"`
Cores int `json:"cores"`
}
type Ram struct {
UsedMB int `json:"usedMb"`
TotalMB int `json:"totalMb"`
UsedPercent int `json:"usedPercent"`
}
type Disk struct {
MountPoint string `json:"mountPoint"`
UsedMB int `json:"usedMb"`
UsedGB int `json:"usedGb"`
TotalMB int `json:"totalMb"`
TotalGB int `json:"totalGb"`
UsedPercent int `json:"usedPercent"`
}
// InitOS 获取OS信息
func InitOS() (o Os) {
o.GOOS = runtime.GOOS
o.NumCPU = runtime.NumCPU()
o.Compiler = runtime.Compiler
o.GoVersion = runtime.Version()
o.NumGoroutine = runtime.NumGoroutine()
return o
}
// InitCPU 获取CPU信息
func InitCPU() (c Cpu, err error) {
if cores, err := cpu.Counts(false); err != nil {
return c, err
} else {
c.Cores = cores
}
if cpus, err := cpu.Percent(time.Duration(200)*time.Millisecond, true); err != nil {
return c, err
} else {
c.Cpus = cpus
}
return c, nil
}
// InitRAM 获取RAM信息
func InitRAM() (r Ram, err error) {
if u, err := mem.VirtualMemory(); err != nil {
return r, err
} else {
r.UsedMB = int(u.Used) / MB
r.TotalMB = int(u.Total) / MB
r.UsedPercent = int(u.UsedPercent)
}
return r, nil
}
// InitDisk 获取硬盘信息
func InitDisk(mountPoints []string) (d []Disk, err error) {
for _, mp := range mountPoints {
if u, err := disk.Usage(mp); err != nil {
return d, err
} else {
d = append(d, Disk{
MountPoint: mp,
UsedMB: int(u.Used) / MB,
UsedGB: int(u.Used) / GB,
TotalMB: int(u.Total) / MB,
TotalGB: int(u.Total) / GB,
UsedPercent: int(u.UsedPercent),
})
}
}
return d, nil
}
+73
View File
@@ -0,0 +1,73 @@
package stacktrace
import (
"regexp"
"strconv"
"strings"
)
// Frame 表示一次栈帧解析结果
type Frame struct {
File string
Line int
Func string
}
var fileLineRe = regexp.MustCompile(`\s*(.+\.go):(\d+)\s*`)
// FindFinalCaller 从 zap 的 entry.Stack 文本中,解析"最终业务调用方"的文件与行号
// 策略:自顶向下解析,优先选择第一条项目代码帧,过滤第三方库/标准库/框架中间件
func FindFinalCaller(stack string) (Frame, bool) {
if stack == "" {
return Frame{}, false
}
lines := strings.Split(stack, "\n")
var currFunc string
for i := 0; i < len(lines); i++ {
line := strings.TrimSpace(lines[i])
if line == "" {
continue
}
if m := fileLineRe.FindStringSubmatch(line); m != nil {
file := m[1]
ln, _ := strconv.Atoi(m[2])
if shouldSkip(file) {
// 跳过此帧,同时重置函数名以避免错误配对
currFunc = ""
continue
}
return Frame{File: file, Line: ln, Func: currFunc}, true
}
// 记录函数名行,下一行通常是文件:行
currFunc = line
}
return Frame{}, false
}
func shouldSkip(file string) bool {
// 第三方库与 Go 模块缓存
if strings.Contains(file, "/go/pkg/mod/") {
return true
}
if strings.Contains(file, "/go.uber.org/") {
return true
}
if strings.Contains(file, "/gorm.io/") {
return true
}
// 标准库
if strings.Contains(file, "/go/go") && strings.Contains(file, "/src/") {
return true
}
// 框架内不需要作为最终调用方的路径
if strings.Contains(file, "/internal/server/middleware/") {
return true
}
if strings.Contains(file, "/internal/server/router/") {
return true
}
if strings.Contains(file, "/pkg/utils/") {
return true
}
return false
}
+34
View File
@@ -0,0 +1,34 @@
package utils
import (
"sync"
)
// SystemEvents 定义系统级事件处理
type SystemEvents struct {
reloadHandlers []func() error
mu sync.RWMutex
}
// GlobalSystemEvents 全局事件管理器
var GlobalSystemEvents = &SystemEvents{}
// RegisterReloadHandler 注册系统重载处理函数
func (e *SystemEvents) RegisterReloadHandler(handler func() error) {
e.mu.Lock()
defer e.mu.Unlock()
e.reloadHandlers = append(e.reloadHandlers, handler)
}
// TriggerReload 触发所有注册的重载处理函数
func (e *SystemEvents) TriggerReload() error {
e.mu.RLock()
defer e.mu.RUnlock()
for _, handler := range e.reloadHandlers {
if err := handler(); err != nil {
return err
}
}
return nil
}
+228 -2
View File
@@ -1,5 +1,22 @@
package utils
import (
"errors"
"reflect"
"regexp"
"strconv"
"strings"
)
// Rules 验证规则类型
type Rules map[string][]string
// RulesMap 规则映射类型
type RulesMap map[string]Rules
// CustomizeMap 自定义规则映射
var CustomizeMap = make(map[string]Rules)
// 预定义验证规则
var (
IdVerify = Rules{"ID": []string{NotEmpty()}}
@@ -19,8 +36,14 @@ var (
SetUserAuthorityVerify = Rules{"AuthorityId": {NotEmpty()}}
)
// Rules 验证规则类型
type Rules map[string][]string
// RegisterRule 注册自定义规则方案建议在路由初始化层即注册
func RegisterRule(key string, rule Rules) (err error) {
if CustomizeMap[key] != nil {
return errors.New(key + "已注册,无法重复注册")
}
CustomizeMap[key] = rule
return nil
}
// NotEmpty 非空验证
func NotEmpty() string {
@@ -76,3 +99,206 @@ func MaxLen(value string) string {
func Regexp(value string) string {
return "regexp=" + value
}
// RegexpMatch 正则校验 校验输入项是否满足正则表达式
func RegexpMatch(rule string) string {
return "regexp=" + rule
}
// Verify 验证结构体字段
func Verify(st interface{}, rules Rules) error {
compareMap := map[string]bool{
"lt": true, "le": true, "eq": true, "ne": true, "ge": true, "gt": true,
}
typ := reflect.TypeOf(st)
val := reflect.ValueOf(st)
// 处理指针类型
if typ.Kind() == reflect.Ptr {
typ = typ.Elem()
val = val.Elem()
}
kd := val.Kind()
if kd != reflect.Struct {
return errors.New("expect struct")
}
num := val.NumField()
for i := 0; i < num; i++ {
tagVal := typ.Field(i)
fieldVal := val.Field(i)
fieldName := tagVal.Name
// 递归处理嵌套结构体
if tagVal.Type.Kind() == reflect.Struct {
if err := Verify(fieldVal.Interface(), rules); err != nil {
return err
}
}
// 检查是否有验证规则
if ruleList, ok := rules[fieldName]; ok {
for _, rule := range ruleList {
switch {
case rule == "notEmpty":
if isBlank(fieldVal) {
return errors.New(fieldName + "值不能为空")
}
case strings.HasPrefix(rule, "regexp="):
if !regexpMatch(strings.Split(rule, "=")[1], fieldVal.String()) {
return errors.New(fieldName + "格式校验不通过")
}
case compareMap[strings.Split(rule, "=")[0]]:
if !compareVerify(fieldVal, rule) {
return errors.New(fieldName + "长度或值不在合法范围," + rule)
}
}
}
}
}
return nil
}
// isBlank 非空校验
func isBlank(value reflect.Value) bool {
switch value.Kind() {
case reflect.String, reflect.Slice:
return value.Len() == 0
case reflect.Bool:
return !value.Bool()
case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
return value.Int() == 0
case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64, reflect.Uintptr:
return value.Uint() == 0
case reflect.Float32, reflect.Float64:
return value.Float() == 0
case reflect.Interface, reflect.Ptr:
return value.IsNil()
}
return reflect.DeepEqual(value.Interface(), reflect.Zero(value.Type()).Interface())
}
// compareVerify 长度和数字的校验方法 根据类型自动校验
func compareVerify(value reflect.Value, VerifyStr string) bool {
switch value.Kind() {
case reflect.String:
return compareInt(len([]rune(value.String())), VerifyStr)
case reflect.Slice, reflect.Array:
return compareInt(value.Len(), VerifyStr)
case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64, reflect.Uintptr:
return compareUint(value.Uint(), VerifyStr)
case reflect.Float32, reflect.Float64:
return compareFloat64(value.Float(), VerifyStr)
case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
return compareInt64(value.Int(), VerifyStr)
default:
return false
}
}
// compareInt 整数比较
func compareInt(value int, VerifyStr string) bool {
VerifyStrArr := strings.Split(VerifyStr, "=")
VInt, VErr := strconv.Atoi(VerifyStrArr[1])
if VErr != nil {
return false
}
switch VerifyStrArr[0] {
case "lt":
return value < VInt
case "le":
return value <= VInt
case "eq":
return value == VInt
case "ne":
return value != VInt
case "ge":
return value >= VInt
case "gt":
return value > VInt
default:
return false
}
}
// compareInt64 int64比较
func compareInt64(value int64, VerifyStr string) bool {
VerifyStrArr := strings.Split(VerifyStr, "=")
VInt, VErr := strconv.ParseInt(VerifyStrArr[1], 10, 64)
if VErr != nil {
return false
}
switch VerifyStrArr[0] {
case "lt":
return value < VInt
case "le":
return value <= VInt
case "eq":
return value == VInt
case "ne":
return value != VInt
case "ge":
return value >= VInt
case "gt":
return value > VInt
default:
return false
}
}
// compareUint uint64比较
func compareUint(value uint64, VerifyStr string) bool {
VerifyStrArr := strings.Split(VerifyStr, "=")
VInt, VErr := strconv.ParseUint(VerifyStrArr[1], 10, 64)
if VErr != nil {
return false
}
switch VerifyStrArr[0] {
case "lt":
return value < VInt
case "le":
return value <= VInt
case "eq":
return value == VInt
case "ne":
return value != VInt
case "ge":
return value >= VInt
case "gt":
return value > VInt
default:
return false
}
}
// compareFloat64 float64比较
func compareFloat64(value float64, VerifyStr string) bool {
VerifyStrArr := strings.Split(VerifyStr, "=")
VFloat, VErr := strconv.ParseFloat(VerifyStrArr[1], 64)
if VErr != nil {
return false
}
switch VerifyStrArr[0] {
case "lt":
return value < VFloat
case "le":
return value <= VFloat
case "eq":
return value == VFloat
case "ne":
return value != VFloat
case "ge":
return value >= VFloat
case "gt":
return value > VFloat
default:
return false
}
}
// regexpMatch 正则匹配
func regexpMatch(rule, matchStr string) bool {
return regexp.MustCompile(rule).MatchString(matchStr)
}