package dao import ( "Cold_Logistic/internal/pkg/common/constant" "Cold_Logistic/internal/pkg/common/global" "context" "gorm.io/gorm" "gorm.io/gorm/clause" "Cold_Logistic/internal/server/infra/models" ) type DataStore struct { db *gorm.DB } type dbOption func(*gorm.DB) *gorm.DB func NewDataStore(db *gorm.DB) *DataStore { return &DataStore{db: db} } type contextTxKey struct{} func (ds *DataStore) InTx(ctx context.Context, fn func(ctx context.Context) error) error { return ds.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error { ctx = context.WithValue(ctx, contextTxKey{}, tx) return fn(ctx) }) } func (ds *DataStore) GetDB(ctx context.Context) *gorm.DB { tx, ok := ctx.Value(contextTxKey{}).(*gorm.DB) if ok { return tx } return ds.db } func (ds *DataStore) withByID(id int) dbOption { return func(db *gorm.DB) *gorm.DB { return db.Where("id = ?", id) } } func (ds *DataStore) withByColumn(column string, val interface{}) dbOption { return func(db *gorm.DB) *gorm.DB { return db.Where(column+" = ?", val) } } func (ds *DataStore) withByColumnInVal(column string, val interface{}) dbOption { return func(db *gorm.DB) *gorm.DB { return db.Where(column+" in (?)", val) } } func (ds *DataStore) withBySelects(filed ...string) dbOption { return func(db *gorm.DB) *gorm.DB { return db.Select(filed) } } func (ds *DataStore) withByNotDeleted() dbOption { return func(db *gorm.DB) *gorm.DB { return db.Where("deleted = ?", models.DeleteNo) } } func (ds *DataStore) withForUpdate() dbOption { return func(db *gorm.DB) *gorm.DB { return db.Clauses(clause.Locking{Strength: "UPDATE"}) } } func (ds *DataStore) optionDB(ctx context.Context, opts ...dbOption) *gorm.DB { db := ds.GetDB(ctx).WithContext(ctx) // db := ds.db.WithContext(ctx) for _, opt := range opts { db = opt(db) } return db } func (ds *DataStore) withAccountDataPermis(ctx context.Context, column string) dbOption { tokenInfo := global.GetTokenInfoFromContext(ctx) if tokenInfo.AccountType == constant.AccountApplet { return func(db *gorm.DB) *gorm.DB { return db.Where(column+" = ?", tokenInfo.AccountId) } } // 默认 return func(db *gorm.DB) *gorm.DB { return db } } // 数据权限 func (ds *DataStore) withCompanyDataPermis(ctx context.Context, column string) dbOption { tokenInfo := global.GetTokenInfoFromContext(ctx) if tokenInfo.AccountType == constant.AccountPlatform { return func(db *gorm.DB) *gorm.DB { return db.Where(column+" = ?", tokenInfo.UsePid) } } // 默认 return func(db *gorm.DB) *gorm.DB { return db } }