package entity

import (
	"context"
	"gorm.io/gorm"
	iContext "yx-dataset-server/app/context"
)


func getDB(ctx context.Context, defDB *gorm.DB) *gorm.DB {
	trans, ok := iContext.FromTrans(ctx)
	if ok {
		db, ok := trans.(*gorm.DB)
		if ok {
			if iContext.FromTransLock(ctx) {
				db = db.Set("gorm:query_option", "FOR UPDATE")
			}
			return db
		}
	}
	return defDB
}

func getDBWithModel(ctx context.Context, defDB *gorm.DB, m interface{}) *gorm.DB {
	return getDB(ctx, defDB).Model(m).WithContext(ctx)
}
