Cold_Logistic/internal/server/infra/dao/a_dataStore.go

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
	}
}