package model import ( "context" "errors" "gorm.io/gorm" ) // AIUserQuotaFindByUser 根据用户ID查询额度 func AIUserQuotaFindByUser(ctx context.Context, db *gorm.DB, userId int64) (*AIUserQuota, error) { var quota AIUserQuota result := db.WithContext(ctx).Where("user_id = ?", userId).First("a) if result.Error != nil { if errors.Is(result.Error, gorm.ErrRecordNotFound) { return nil, ErrNotFound } return nil, result.Error } return "a, nil } // AIUserQuotaEnsure 查找或创建用户额度记录 func AIUserQuotaEnsure(ctx context.Context, db *gorm.DB, userId int64) (*AIUserQuota, error) { var quota AIUserQuota result := db.WithContext(ctx).Where("user_id = ?", userId).First("a) if result.Error != nil { if errors.Is(result.Error, gorm.ErrRecordNotFound) { quota = AIUserQuota{UserId: userId} if err := db.WithContext(ctx).Create("a).Error; err != nil { return nil, err } return "a, nil } return nil, result.Error } return "a, nil } // AIUserQuotaFreeze 冻结额度(原子操作:balance -= amount, frozen_amount += amount) func AIUserQuotaFreeze(ctx context.Context, db *gorm.DB, userId int64, amount float64) error { result := db.WithContext(ctx).Model(&AIUserQuota{}). Where("user_id = ? AND balance >= ?", userId, amount). Updates(map[string]interface{}{ "balance": gorm.Expr("balance - ?", amount), "frozen_amount": gorm.Expr("frozen_amount + ?", amount), }) if result.Error != nil { return result.Error } if result.RowsAffected == 0 { return errors.New("insufficient balance") } return nil } // AIUserQuotaSettle 结算额度(原子操作:frozen_amount -= frozenAmount, total_consumed += actualCost, balance += refund) func AIUserQuotaSettle(ctx context.Context, db *gorm.DB, userId int64, frozenAmount, actualCost float64) error { refund := frozenAmount - actualCost result := db.WithContext(ctx).Model(&AIUserQuota{}). Where("user_id = ?", userId). Updates(map[string]interface{}{ "frozen_amount": gorm.Expr("frozen_amount - ?", frozenAmount), "total_consumed": gorm.Expr("total_consumed + ?", actualCost), "balance": gorm.Expr("balance + ?", refund), }) if result.Error != nil { return result.Error } if result.RowsAffected == 0 { return ErrNotFound } return nil } // AIUserQuotaFindList 分页查询所有用户额度(admin用) func AIUserQuotaFindList(ctx context.Context, db *gorm.DB, page, pageSize int64) ([]AIUserQuota, int64, error) { var list []AIUserQuota var total int64 query := db.WithContext(ctx).Model(&AIUserQuota{}) if err := query.Count(&total).Error; err != nil { return nil, 0, err } offset := (page - 1) * pageSize if offset < 0 { offset = 0 } err := query.Order("user_id ASC").Offset(int(offset)).Limit(int(pageSize)).Find(&list).Error return list, total, err } // AIUserQuotaRecharge 充值(原子操作:balance += amount, total_recharged += amount) func AIUserQuotaRecharge(ctx context.Context, db *gorm.DB, userId int64, amount float64) error { // Ensure record exists _, err := AIUserQuotaEnsure(ctx, db, userId) if err != nil { return err } return db.WithContext(ctx).Model(&AIUserQuota{}). Where("user_id = ?", userId). Updates(map[string]interface{}{ "balance": gorm.Expr("balance + ?", amount), "total_recharged": gorm.Expr("total_recharged + ?", amount), }).Error } // AIUserQuotaUnfreeze 解冻额度(原子操作:frozen_amount -= amount, balance += amount) func AIUserQuotaUnfreeze(ctx context.Context, db *gorm.DB, userId int64, amount float64) error { result := db.WithContext(ctx).Model(&AIUserQuota{}). Where("user_id = ?", userId). Updates(map[string]interface{}{ "frozen_amount": gorm.Expr("frozen_amount - ?", amount), "balance": gorm.Expr("balance + ?", amount), }) if result.Error != nil { return result.Error } if result.RowsAffected == 0 { return ErrNotFound } return nil }