feat: 添加BusinessDB字段到PluginInitializeGorm结构体，并增加相关测试用例

pixelmaxQm · pixelmaxQm · commit 15fdb73b53c8 · 2026-04-07T12:15:59.000+08:00
diff --git a/server/service/system/auto_code_package.go b/server/service/system/auto_code_package.go
@@ -571,6 +571,7 @@ func (s *autoCodePackage) templates(ctx context.Context, entity model.SysAutoCod
 								Type:        ast.TypePluginInitializeGorm,
 								Path:        filepath.Join(global.GVA_CONFIG.AutoCode.Root, global.GVA_CONFIG.AutoCode.Server, "plugin", entity.PackageName, secondDirs[j].Name(), strings.TrimSuffix(threeDirs[k].Name(), ext)),
 								ImportPath:  fmt.Sprintf(`"%s/plugin/%s/model"`, global.GVA_CONFIG.AutoCode.Module, entity.PackageName),
+								Business:    info.BusinessDB,
 								StructName:  info.StructName,
 								PackageName: "model",
 								IsNew:       true,
diff --git a/server/service/system/auto_code_plugin_gorm_template_test.go b/server/service/system/auto_code_plugin_gorm_template_test.go
@@ -0,0 +1,69 @@
+package system
+
+import (
+	"context"
+	"path/filepath"
+	"reflect"
+	"runtime"
+	"testing"
+
+	"github.com/flipped-aurora/gin-vue-admin/server/global"
+	model "github.com/flipped-aurora/gin-vue-admin/server/model/system"
+	"github.com/flipped-aurora/gin-vue-admin/server/model/system/request"
+	utilsAst "github.com/flipped-aurora/gin-vue-admin/server/utils/ast"
+)
+
+func TestPluginInitializeGormInjectionCarriesBusinessDB(t *testing.T) {
+	_, currentFile, _, ok := runtime.Caller(0)
+	if !ok {
+		t.Fatal("runtime.Caller() failed")
+	}
+	repoRoot := filepath.Clean(filepath.Join(filepath.Dir(currentFile), "..", "..", ".."))
+
+	oldRoot := global.GVA_CONFIG.AutoCode.Root
+	oldServer := global.GVA_CONFIG.AutoCode.Server
+	oldModule := global.GVA_CONFIG.AutoCode.Module
+	global.GVA_CONFIG.AutoCode.Root = repoRoot
+	global.GVA_CONFIG.AutoCode.Server = "server"
+	global.GVA_CONFIG.AutoCode.Module = "github.com/flipped-aurora/gin-vue-admin/server"
+	defer func() {
+		global.GVA_CONFIG.AutoCode.Root = oldRoot
+		global.GVA_CONFIG.AutoCode.Server = oldServer
+		global.GVA_CONFIG.AutoCode.Module = oldModule
+	}()
+
+	info := request.AutoCode{
+		Package:         "demoPlugin",
+		PackageName:     "demo",
+		HumpPackageName: "demo",
+		StructName:      "Demo",
+		Abbreviation:    "demo",
+		BusinessDB:      "bizdb",
+		GenerateServer:  true,
+	}
+	entity := model.SysAutoCodePackage{
+		Template:    "plugin",
+		PackageName: info.Package,
+	}
+
+	_, asts, _, err := AutoCodePackage.templates(context.Background(), entity, info, false)
+	if err != nil {
+		t.Fatalf("templates() error = %v", err)
+	}
+
+	var pluginInitializeGorm *utilsAst.PluginInitializeGorm
+	for _, injection := range asts {
+		if candidate, ok := injection.(*utilsAst.PluginInitializeGorm); ok {
+			pluginInitializeGorm = candidate
+			break
+		}
+	}
+	if pluginInitializeGorm == nil {
+		t.Fatal("expected plugin initialize gorm injection")
+	}
+
+	businessField := reflect.ValueOf(pluginInitializeGorm).Elem().FieldByName("Business")
+	if !businessField.IsValid() || businessField.String() != info.BusinessDB {
+		t.Fatalf("expected PluginInitializeGorm.Business = %q, got %v", info.BusinessDB, businessField)
+	}
+}
diff --git a/server/utils/ast/plugin_initialize_gorm.go b/server/utils/ast/plugin_initialize_gorm.go
@@ -1,7 +1,12 @@
 package ast
 
 import (
+	"bytes"
+	"fmt"
 	"go/ast"
+	"go/format"
+	"go/parser"
+	"go/token"
 	"io"
 )
 
@@ -11,9 +16,10 @@ type PluginInitializeGorm struct {
 	Path         string // 文件路径
 	ImportPath   string // 导包路径
 	RelativePath string // 相对路径
+	Business     string // 业务库
 	StructName   string // 结构体名称
 	PackageName  string // 包名
-	IsNew        bool   // 是否使用new关键字 true: new(PackageName.StructName) false: &PackageName.StructName{}
+	IsNew        bool   // 是否使用 new 关键字 true: new(PackageName.StructName) false: &PackageName.StructName{}
 }
 
 func (a *PluginInitializeGorm) Parse(filename string, writer io.Writer) (file *ast.File, err error) {
@@ -44,7 +50,7 @@ func (a *PluginInitializeGorm) Rollback(file *ast.File) error {
 		if len(callExpr.Args) <= 1 {
 			needRollBackImport = true
 		}
-		// 删除指定的参数
+		// 删除指定参数
 		for i, arg := range callExpr.Args {
 			compLit, cok := arg.(*ast.CompositeLit)
 			if !cok {
@@ -76,30 +82,44 @@ func (a *PluginInitializeGorm) Rollback(file *ast.File) error {
 
 func (a *PluginInitializeGorm) Injection(file *ast.File) error {
 	_ = NewImport(a.ImportPath).Injection(file)
-	var call *ast.CallExpr
+
+	var targetCall *ast.CallExpr
 	ast.Inspect(file, func(n ast.Node) bool {
 		callExpr, ok := n.(*ast.CallExpr)
 		if !ok {
 			return true
 		}
 
 		selExpr, ok := callExpr.Fun.(*ast.SelectorExpr)
-		if ok && selExpr.Sel.Name == "AutoMigrate" {
-			call = callExpr
+		if !ok || selExpr.Sel.Name != "AutoMigrate" {
+			return true
+		}
+
+		if a.isTargetAutoMigrateCall(callExpr) {
+			targetCall = callExpr
 			return false
 		}
 
 		return true
 	})
 
-	arg := &ast.CompositeLit{
+	if targetCall == nil {
+		targetCall = a.appendAutoMigrateBlock(file)
+	}
+	if targetCall == nil {
+		return nil
+	}
+
+	if a.hasModelArg(targetCall) {
+		return nil
+	}
+
+	targetCall.Args = append(targetCall.Args, &ast.CompositeLit{
 		Type: &ast.SelectorExpr{
 			X:   &ast.Ident{Name: a.PackageName},
 			Sel: &ast.Ident{Name: a.StructName},
 		},
-	}
-
-	call.Args = append(call.Args, arg)
+	})
 	return nil
 }
 
@@ -109,3 +129,95 @@ func (a *PluginInitializeGorm) Format(filename string, writer io.Writer, file *a
 	}
 	return a.Base.Format(filename, writer, file)
 }
+
+func (a *PluginInitializeGorm) isTargetAutoMigrateCall(callExpr *ast.CallExpr) bool {
+	selExpr, ok := callExpr.Fun.(*ast.SelectorExpr)
+	if !ok || selExpr.Sel.Name != "AutoMigrate" {
+		return false
+	}
+	return exprString(selExpr.X) == exprString(a.autoMigrateReceiverExpr())
+}
+
+func (a *PluginInitializeGorm) appendAutoMigrateBlock(file *ast.File) *ast.CallExpr {
+	gormFunc := FindFunction(file, "Gorm")
+	if gormFunc == nil || gormFunc.Body == nil {
+		return nil
+	}
+
+	src := fmt.Sprintf(`package placeholder
+func Gorm() {
+	if err = %s.AutoMigrate(); err != nil {
+		err = errors.Wrap(err, "注册表失败!")
+		zap.L().Error(fmt.Sprintf("%%+v", err))
+	}
+}
+`, exprString(a.autoMigrateReceiverExpr()))
+
+	parsed, err := parser.ParseFile(token.NewFileSet(), "", src, 0)
+	if err != nil || len(parsed.Decls) == 0 {
+		return nil
+	}
+
+	stmt := parsed.Decls[0].(*ast.FuncDecl).Body.List[0].(*ast.IfStmt)
+	clearPosition(stmt)
+	gormFunc.Body.List = append(gormFunc.Body.List, stmt)
+
+	assignStmt := stmt.Init.(*ast.AssignStmt)
+	callExpr := assignStmt.Rhs[0].(*ast.CallExpr)
+	return callExpr
+}
+
+func (a *PluginInitializeGorm) autoMigrateReceiverExpr() ast.Expr {
+	return &ast.CallExpr{
+		Fun: &ast.SelectorExpr{
+			X:   a.dbExpr(),
+			Sel: &ast.Ident{Name: "WithContext"},
+		},
+		Args: []ast.Expr{&ast.Ident{Name: "ctx"}},
+	}
+}
+
+func (a *PluginInitializeGorm) dbExpr() ast.Expr {
+	if a.Business == "" {
+		return &ast.SelectorExpr{
+			X:   &ast.Ident{Name: "global"},
+			Sel: &ast.Ident{Name: "GVA_DB"},
+		}
+	}
+	return &ast.CallExpr{
+		Fun: &ast.SelectorExpr{
+			X:   &ast.Ident{Name: "global"},
+			Sel: &ast.Ident{Name: "MustGetGlobalDBByDBName"},
+		},
+		Args: []ast.Expr{
+			&ast.BasicLit{
+				Kind:  token.STRING,
+				Value: fmt.Sprintf("\"%s\"", a.Business),
+			},
+		},
+	}
+}
+
+func (a *PluginInitializeGorm) hasModelArg(callExpr *ast.CallExpr) bool {
+	for _, arg := range callExpr.Args {
+		compositeLit, ok := arg.(*ast.CompositeLit)
+		if !ok {
+			continue
+		}
+		selectorExpr, ok := compositeLit.Type.(*ast.SelectorExpr)
+		if !ok {
+			continue
+		}
+		packageIdent, ok := selectorExpr.X.(*ast.Ident)
+		if ok && packageIdent.Name == a.PackageName && selectorExpr.Sel.Name == a.StructName {
+			return true
+		}
+	}
+	return false
+}
+
+func exprString(expr ast.Expr) string {
+	var buffer bytes.Buffer
+	_ = format.Node(&buffer, token.NewFileSet(), expr)
+	return buffer.String()
+}
diff --git a/server/utils/ast/plugin_initialize_gorm_business_test.go b/server/utils/ast/plugin_initialize_gorm_business_test.go
@@ -0,0 +1,79 @@
+package ast
+
+import (
+	"os"
+	"path/filepath"
+	"reflect"
+	"strings"
+	"testing"
+)
+
+func TestPluginInitializeGormInjectionUsesBusinessDB(t *testing.T) {
+	const source = `package initialize
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/flipped-aurora/gin-vue-admin/server/global"
+	"github.com/pkg/errors"
+	"go.uber.org/zap"
+)
+
+func Gorm(ctx context.Context) {
+	err := global.GVA_DB.WithContext(ctx).AutoMigrate()
+	if err != nil {
+		err = errors.Wrap(err, "注册表失败!")
+		zap.L().Error(fmt.Sprintf("%+v", err))
+	}
+}
+`
+
+	dir := t.TempDir()
+	path := filepath.Join(dir, "gorm.go")
+	if err := os.WriteFile(path, []byte(source), 0o666); err != nil {
+		t.Fatalf("WriteFile() error = %v", err)
+	}
+
+	injection := &PluginInitializeGorm{
+		Type:        TypePluginInitializeGorm,
+		Path:        path,
+		ImportPath:  `"github.com/flipped-aurora/gin-vue-admin/server/plugin/demo/model"`,
+		StructName:  "Demo",
+		PackageName: "model",
+		IsNew:       true,
+	}
+
+	businessField := reflect.ValueOf(injection).Elem().FieldByName("Business")
+	if !businessField.IsValid() {
+		t.Fatal("expected PluginInitializeGorm.Business field")
+	}
+	businessField.SetString("bizdb")
+
+	file, err := injection.Parse(path, nil)
+	if err != nil {
+		t.Fatalf("Parse() error = %v", err)
+	}
+	if err := injection.Injection(file); err != nil {
+		t.Fatalf("Injection() error = %v", err)
+	}
+	if err := injection.Format(path, nil, file); err != nil {
+		t.Fatalf("Format() error = %v", err)
+	}
+
+	content, err := os.ReadFile(path)
+	if err != nil {
+		t.Fatalf("ReadFile() error = %v", err)
+	}
+
+	got := string(content)
+	if !strings.Contains(got, "global.GVA_DB.WithContext(ctx).AutoMigrate()") {
+		t.Fatalf("expected default gorm block to remain, got:\n%s", got)
+	}
+	if !strings.Contains(got, `global.MustGetGlobalDBByDBName("bizdb").WithContext(ctx).AutoMigrate(`) {
+		t.Fatalf("expected gorm injection to use business db, got:\n%s", got)
+	}
+	if !strings.Contains(got, "model.Demo{}") {
+		t.Fatalf("expected model injection, got:\n%s", got)
+	}
+}
