70 lines
1.5 KiB
Go
70 lines
1.5 KiB
Go
package store
|
|
|
|
import (
|
|
"context"
|
|
"sync"
|
|
|
|
"management/internal/erpserver/store/system"
|
|
|
|
"gorm.io/gorm"
|
|
)
|
|
|
|
var (
|
|
once sync.Once
|
|
// 全局变量，方便其它包直接调用已初始化好的 datastore 实例.
|
|
engine *datastore
|
|
)
|
|
|
|
// IStore 定义了 Store 层需要实现的方法.
|
|
type IStore interface {
|
|
DB(ctx context.Context) *gorm.DB
|
|
TX(ctx context.Context, fn func(ctx context.Context) error) error
|
|
|
|
User() system.UserStore
|
|
}
|
|
|
|
// transactionKey 用于在 context.Context 中存储事务上下文的键.
|
|
type transactionKey struct{}
|
|
|
|
// datastore 是 IStore 的具体实现.
|
|
type datastore struct {
|
|
core *gorm.DB
|
|
}
|
|
|
|
// 确保 datastore 实现了 IStore 接口.
|
|
var _ IStore = (*datastore)(nil)
|
|
|
|
// NewStore 创建一个 IStore 类型的实例.
|
|
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)
|
|
},
|
|
)
|
|
}
|
|
|
|
// Users 返回一个实现了 UserStore 接口的实例.
|
|
func (store *datastore) User() system.UserStore {
|
|
return system.NewUserStore(store.core)
|
|
}
|