feat: able to set initial quota for new user (close #22)

model/user.go

@@ -2,6 +2,7 @@ package model
 
 import (
 	"errors"
+	"gorm.io/gorm"
 	"one-api/common"
 	"strings"
 )
@@ -21,6 +22,7 @@ type User struct {
 	VerificationCode string `json:"verification_code" gorm:"-:all"` // this field is only for Email verification, don't save it to database!
 	Balance          int    `json:"balance" gorm:"type:int;default:0"`
 	AccessToken      string `json:"access_token" gorm:"column:access_token;uniqueIndex"` // this token is for system management
+	Quota            int    `json:"quota" gorm:"type:int;default:0"`
 }
 
 func GetMaxUserId() int {
@@ -69,6 +71,8 @@ func (user *User) Insert() error {
 			return err
 		}
 	}
+	user.Quota = common.QuotaForNewUser
+	user.AccessToken = common.GetUUID()
 	err = DB.Create(user).Error
 	return err
 }
@@ -202,3 +206,13 @@ func ValidateAccessToken(token string) (user *User) {
 	}
 	return nil
 }
+
+func GetUserQuota(id int) (quota int, err error) {
+	err = DB.Model(&User{}).Where("id = ?", id).Select("quota").Find(&quota).Error
+	return quota, err
+}
+
+func DecreaseUserQuota(id int, quota int) (err error) {
+	err = DB.Model(&User{}).Where("id = ?", id).Update("quota", gorm.Expr("quota - ?", quota)).Error
+	return err
+}