feat: orm 支持 gorm 事务封装

Author: sendya

kratos/orm/gorm_logger.go

@@ -93,7 +93,7 @@ func (gl *gormLogger) Trace(ctx context.Context, begin time.Time, fc func() (str
 	timeUsed := float64(elapsed.Nanoseconds()) / 1e6
 
 	fields := make([]interface{}, 0)
-	fields = append(fields, "timeUsedMicroseconds", timeUsed)
+	fields = append(fields, "timeUsed", timeUsed)
 
 	sql, rows := fc()
 
@@ -108,7 +108,7 @@ func (gl *gormLogger) Trace(ctx context.Context, begin time.Time, fc func() (str
 	case elapsed > gl.SlowThreshold && gl.SlowThreshold != 0:
 		fields = append(fields, "sql", sql)
 		fields = append(fields, "rows", rows)
-		fields = append(fields, "slow-elapsed", gl.SlowThreshold)
+		fields = append(fields, "slowElapsed", gl.SlowThreshold)
 		gl.dbLog.Warnw(fields...)
 	// normal
 	default:

kratos/orm/gorm_transaction_test.go

@@ -0,0 +1,67 @@
+package orm_test
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/glebarez/sqlite"
+	"github.com/go-kratos/kratos/v2/log"
+	"gorm.io/gorm"
+
+	"github.com/omalloc/contrib/kratos/orm"
+)
+
+var cfg = `
+server:
+  http:
+    addr: 0.0.0.0:8000
+    timeout: 1s
+  grpc:
+    addr: 0.0.0.0:9000
+    timeout: 1s
+data:
+  database:
+    driver: sqlite
+    source: file:test.db?cache=shared&mode=memory
+  redis:
+    addr: 127.0.0.1:6379
+    read_timeout: 0.2s
+    write_timeout: 0.2s
+`
+
+type Data struct {
+	db *gorm.DB
+}
+
+func newData() *Data {
+	db, err := orm.New(
+		orm.WithDriver(sqlite.Open("file:test.db?cache=shared&mode=memory&charset=utf8mb4&parseTime=true&loc=Local")),
+		orm.WithTracingOpts(orm.WithDatabaseName("test")),
+		orm.WithLogger(
+			orm.WithDebug(),
+			orm.WIthSlowThreshold(time.Second*2),
+			orm.WithSkipCallerLookup(true),
+			orm.WithSkipErrRecordNotFound(true),
+			orm.WithLogHelper(log.NewFilter(log.GetLogger(), log.FilterLevel(log.LevelDebug))),
+		),
+	)
+	if err != nil {
+		panic(err)
+	}
+
+	return &Data{
+		db: db,
+	}
+}
+
+func TestTransaction(t *testing.T) {
+	data := newData()
+	data.db.AutoMigrate(&User{})
+
+	txm := orm.NewTransactionManager(data.db)
+
+	txm.WithContext(context.TODO(), func(ctx context.Context, tx *gorm.DB) error {
+		return tx.Model(&User{}).Create(&User{Name: "test1"}).Error
+	})
+}

kratos/orm/transaction.go

@@ -0,0 +1,34 @@
+package orm
+
+import (
+	"context"
+
+	"gorm.io/gorm"
+)
+
+// txContextKey gorm database transaction context key
+type txContextKey struct{}
+
+type Transaction interface {
+	WithContext(context.Context, func(ctx context.Context, tx *gorm.DB) error) error
+}
+
+type transactionManager struct {
+	db *gorm.DB
+}
+
+func NewTransactionManager(dataSource *gorm.DB) Transaction {
+	return &transactionManager{
+		db: dataSource,
+	}
+}
+
+func (tm *transactionManager) WithContext(ctx context.Context, fn func(ctx context.Context, tx *gorm.DB) error) error {
+	if tx, ok := ctx.Value(txContextKey{}).(*gorm.DB); ok {
+		return fn(ctx, tx)
+	}
+
+	return tm.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
+		return fn(context.WithValue(ctx, txContextKey{}, tx), tx)
+	})
+}