60 lines
1.2 KiB
Go
60 lines
1.2 KiB
Go
package repository
|
|
|
|
import (
|
|
"context"
|
|
"sync"
|
|
|
|
"gorm.io/gorm"
|
|
)
|
|
|
|
var (
|
|
once sync.Once
|
|
// 全局变量，方便其它包直接调用已初始化好的 datastore 实例.
|
|
engine *datastore
|
|
)
|
|
|
|
type Store interface {
|
|
DB(ctx context.Context) *gorm.DB
|
|
TX(ctx context.Context, fn func(ctx context.Context) error) error
|
|
}
|
|
|
|
// transactionKey 用于在 context.Context 中存储事务上下文的键.
|
|
type transactionKey struct{}
|
|
|
|
// datastore 是 Storer 的具体实现.
|
|
type datastore struct {
|
|
core *gorm.DB
|
|
}
|
|
|
|
// 确保 datastore 实现了 Storer 接口.
|
|
var _ Store = (*datastore)(nil)
|
|
|
|
// NewStore 创建一个 Storer 类型的实例.
|
|
func NewStore(db *gorm.DB) *datastore {
|
|
// 确保 engine 只被初始化一次
|
|
once.Do(func() {
|
|
engine = &datastore{db}
|
|
})
|
|
|
|
return engine
|
|
}
|
|
|
|
func (store *datastore) DB(ctx context.Context) *gorm.DB {
|
|
db := store.core
|
|
// 从上下文中提取事务实例
|
|
if tx, ok := ctx.Value(transactionKey{}).(*gorm.DB); ok {
|
|
db = tx
|
|
}
|
|
|
|
return db
|
|
}
|
|
|
|
func (store *datastore) TX(ctx context.Context, fn func(ctx context.Context) error) error {
|
|
return store.core.WithContext(ctx).Transaction(
|
|
func(tx *gorm.DB) error {
|
|
ctx = context.WithValue(ctx, transactionKey{}, tx)
|
|
return fn(ctx)
|
|
},
|
|
)
|
|
}
|