feat: 重构事务管理器，支持通过 DataSourceManager 获取数据源

Author: sendya

kratos/orm/gorm_transaction_test.go

@@ -37,13 +37,21 @@ func newData() *Data {
 	}
 }
 
+func (d *Data) GetDataSource() *gorm.DB {
+	return d.db
+}
+
 func TestTransaction(t *testing.T) {
 	data := newData()
 	data.db.AutoMigrate(&User{})
 
-	txm := orm.NewTransactionManager(data.db)
+	txm := orm.NewTransactionManager(data)
+
+	txm.Transaction(context.Background(), func(ctx context.Context) error {
+		// do something
+
+		// do something ..
 
-	txm.WithContext(context.TODO(), func(ctx context.Context, tx *gorm.DB) error {
-		return tx.Model(&User{}).Create(&User{Name: "test1"}).Error
+		return nil
 	})
 }

kratos/orm/transaction.go

@@ -9,26 +9,51 @@ import (
 // txContextKey gorm database transaction context key
 type txContextKey struct{}
 
+type DataSourceManager interface {
+	GetDataSource() *gorm.DB
+}
+
 type Transaction interface {
-	WithContext(context.Context, func(ctx context.Context, tx *gorm.DB) error) error
+	Transaction(context.Context, func(ctx context.Context) error) error
+	WithContext(context.Context) *gorm.DB
 }
 
 type transactionManager struct {
-	db *gorm.DB
+	dsm DataSourceManager
 }
 
-func NewTransactionManager(dataSource *gorm.DB) Transaction {
+func NewTransactionManager(dsm DataSourceManager) Transaction {
 	return &transactionManager{
-		db: dataSource,
+		dsm: dsm,
 	}
 }
 
-func (tm *transactionManager) WithContext(ctx context.Context, fn func(ctx context.Context, tx *gorm.DB) error) error {
+func (tm *transactionManager) WithContext(ctx context.Context) *gorm.DB {
+	if ctx == nil {
+		return tm.dsm.GetDataSource()
+	}
+
 	if tx, ok := ctx.Value(txContextKey{}).(*gorm.DB); ok {
-		return fn(ctx, tx)
+		return tx
+	}
+
+	return tm.dsm.GetDataSource().WithContext(ctx)
+}
+
+func (tm *transactionManager) Transaction(ctx context.Context, fn func(ctx context.Context) error) error {
+	if ctx == nil {
+		return tm.with(context.Background(), fn)
+	}
+
+	if _, ok := ctx.Value(txContextKey{}).(*gorm.DB); ok {
+		return fn(ctx)
 	}
 
-	return tm.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
-		return fn(context.WithValue(ctx, txContextKey{}, tx), tx)
+	return tm.with(ctx, fn)
+}
+
+func (tm *transactionManager) with(ctx context.Context, fn func(ctx context.Context) error) error {
+	return tm.dsm.GetDataSource().Transaction(func(tx *gorm.DB) error {
+		return fn(context.WithValue(ctx, txContextKey{}, tx))
 	})
 }