user.go 23 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822
  1. package model
  2. import (
  3. "encoding/json"
  4. "errors"
  5. "fmt"
  6. "one-api/common"
  7. "strconv"
  8. "strings"
  9. "github.com/bytedance/gopkg/util/gopool"
  10. "gorm.io/gorm"
  11. )
  12. // User if you add sensitive fields, don't forget to clean them in setupLogin function.
  13. // Otherwise, the sensitive information will be saved on local storage in plain text!
  14. type User struct {
  15. Id int `json:"id"`
  16. Username string `json:"username" gorm:"unique;index" validate:"max=12"`
  17. Password string `json:"password" gorm:"not null;" validate:"min=8,max=20"`
  18. OriginalPassword string `json:"original_password" gorm:"-:all"` // this field is only for Password change verification, don't save it to database!
  19. DisplayName string `json:"display_name" gorm:"index" validate:"max=20"`
  20. Role int `json:"role" gorm:"type:int;default:1"` // admin, common
  21. Status int `json:"status" gorm:"type:int;default:1"` // enabled, disabled
  22. Email string `json:"email" gorm:"index" validate:"max=50"`
  23. GitHubId string `json:"github_id" gorm:"column:github_id;index"`
  24. OidcId string `json:"oidc_id" gorm:"column:oidc_id;index"`
  25. WeChatId string `json:"wechat_id" gorm:"column:wechat_id;index"`
  26. TelegramId string `json:"telegram_id" gorm:"column:telegram_id;index"`
  27. VerificationCode string `json:"verification_code" gorm:"-:all"` // this field is only for Email verification, don't save it to database!
  28. AccessToken *string `json:"access_token" gorm:"type:char(32);column:access_token;uniqueIndex"` // this token is for system management
  29. Quota int `json:"quota" gorm:"type:int;default:0"`
  30. UsedQuota int `json:"used_quota" gorm:"type:int;default:0;column:used_quota"` // used quota
  31. RequestCount int `json:"request_count" gorm:"type:int;default:0;"` // request number
  32. Group string `json:"group" gorm:"type:varchar(64);default:'default'"`
  33. AffCode string `json:"aff_code" gorm:"type:varchar(32);column:aff_code;uniqueIndex"`
  34. AffCount int `json:"aff_count" gorm:"type:int;default:0;column:aff_count"`
  35. AffQuota int `json:"aff_quota" gorm:"type:int;default:0;column:aff_quota"` // 邀请剩余额度
  36. AffHistoryQuota int `json:"aff_history_quota" gorm:"type:int;default:0;column:aff_history"` // 邀请历史额度
  37. InviterId int `json:"inviter_id" gorm:"type:int;column:inviter_id;index"`
  38. DeletedAt gorm.DeletedAt `gorm:"index"`
  39. LinuxDOId string `json:"linux_do_id" gorm:"column:linux_do_id;index"`
  40. Setting string `json:"setting" gorm:"type:text;column:setting"`
  41. Remark string `json:"remark,omitempty" gorm:"type:varchar(255)" validate:"max=255"`
  42. }
  43. func (user *User) ToBaseUser() *UserBase {
  44. cache := &UserBase{
  45. Id: user.Id,
  46. Group: user.Group,
  47. Quota: user.Quota,
  48. Status: user.Status,
  49. Username: user.Username,
  50. Setting: user.Setting,
  51. Email: user.Email,
  52. }
  53. return cache
  54. }
  55. func (user *User) GetAccessToken() string {
  56. if user.AccessToken == nil {
  57. return ""
  58. }
  59. return *user.AccessToken
  60. }
  61. func (user *User) SetAccessToken(token string) {
  62. user.AccessToken = &token
  63. }
  64. func (user *User) GetSetting() map[string]interface{} {
  65. if user.Setting == "" {
  66. return nil
  67. }
  68. return common.StrToMap(user.Setting)
  69. }
  70. func (user *User) SetSetting(setting map[string]interface{}) {
  71. settingBytes, err := json.Marshal(setting)
  72. if err != nil {
  73. common.SysError("failed to marshal setting: " + err.Error())
  74. return
  75. }
  76. user.Setting = string(settingBytes)
  77. }
  78. // CheckUserExistOrDeleted check if user exist or deleted, if not exist, return false, nil, if deleted or exist, return true, nil
  79. func CheckUserExistOrDeleted(username string, email string) (bool, error) {
  80. var user User
  81. // err := DB.Unscoped().First(&user, "username = ? or email = ?", username, email).Error
  82. // check email if empty
  83. var err error
  84. if email == "" {
  85. err = DB.Unscoped().First(&user, "username = ?", username).Error
  86. } else {
  87. err = DB.Unscoped().First(&user, "username = ? or email = ?", username, email).Error
  88. }
  89. if err != nil {
  90. if errors.Is(err, gorm.ErrRecordNotFound) {
  91. // not exist, return false, nil
  92. return false, nil
  93. }
  94. // other error, return false, err
  95. return false, err
  96. }
  97. // exist, return true, nil
  98. return true, nil
  99. }
  100. func GetMaxUserId() int {
  101. var user User
  102. DB.Unscoped().Last(&user)
  103. return user.Id
  104. }
  105. func GetAllUsers(startIdx int, num int) (users []*User, total int64, err error) {
  106. // Start transaction
  107. tx := DB.Begin()
  108. if tx.Error != nil {
  109. return nil, 0, tx.Error
  110. }
  111. defer func() {
  112. if r := recover(); r != nil {
  113. tx.Rollback()
  114. }
  115. }()
  116. // Get total count within transaction
  117. err = tx.Unscoped().Model(&User{}).Count(&total).Error
  118. if err != nil {
  119. tx.Rollback()
  120. return nil, 0, err
  121. }
  122. // Get paginated users within same transaction
  123. err = tx.Unscoped().Order("id desc").Limit(num).Offset(startIdx).Omit("password").Find(&users).Error
  124. if err != nil {
  125. tx.Rollback()
  126. return nil, 0, err
  127. }
  128. // Commit transaction
  129. if err = tx.Commit().Error; err != nil {
  130. return nil, 0, err
  131. }
  132. return users, total, nil
  133. }
  134. func SearchUsers(keyword string, group string, startIdx int, num int) ([]*User, int64, error) {
  135. var users []*User
  136. var total int64
  137. var err error
  138. // 开始事务
  139. tx := DB.Begin()
  140. if tx.Error != nil {
  141. return nil, 0, tx.Error
  142. }
  143. defer func() {
  144. if r := recover(); r != nil {
  145. tx.Rollback()
  146. }
  147. }()
  148. // 构建基础查询
  149. query := tx.Unscoped().Model(&User{})
  150. // 构建搜索条件
  151. likeCondition := "username LIKE ? OR email LIKE ? OR display_name LIKE ?"
  152. // 尝试将关键字转换为整数ID
  153. keywordInt, err := strconv.Atoi(keyword)
  154. if err == nil {
  155. // 如果是数字,同时搜索ID和其他字段
  156. likeCondition = "id = ? OR " + likeCondition
  157. if group != "" {
  158. query = query.Where("("+likeCondition+") AND "+commonGroupCol+" = ?",
  159. keywordInt, "%"+keyword+"%", "%"+keyword+"%", "%"+keyword+"%", group)
  160. } else {
  161. query = query.Where(likeCondition,
  162. keywordInt, "%"+keyword+"%", "%"+keyword+"%", "%"+keyword+"%")
  163. }
  164. } else {
  165. // 非数字关键字,只搜索字符串字段
  166. if group != "" {
  167. query = query.Where("("+likeCondition+") AND "+commonGroupCol+" = ?",
  168. "%"+keyword+"%", "%"+keyword+"%", "%"+keyword+"%", group)
  169. } else {
  170. query = query.Where(likeCondition,
  171. "%"+keyword+"%", "%"+keyword+"%", "%"+keyword+"%")
  172. }
  173. }
  174. // 获取总数
  175. err = query.Count(&total).Error
  176. if err != nil {
  177. tx.Rollback()
  178. return nil, 0, err
  179. }
  180. // 获取分页数据
  181. err = query.Omit("password").Order("id desc").Limit(num).Offset(startIdx).Find(&users).Error
  182. if err != nil {
  183. tx.Rollback()
  184. return nil, 0, err
  185. }
  186. // 提交事务
  187. if err = tx.Commit().Error; err != nil {
  188. return nil, 0, err
  189. }
  190. return users, total, nil
  191. }
  192. func GetUserById(id int, selectAll bool) (*User, error) {
  193. if id == 0 {
  194. return nil, errors.New("id 为空!")
  195. }
  196. user := User{Id: id}
  197. var err error = nil
  198. if selectAll {
  199. err = DB.First(&user, "id = ?", id).Error
  200. } else {
  201. err = DB.Omit("password").First(&user, "id = ?", id).Error
  202. }
  203. return &user, err
  204. }
  205. func GetUserIdByAffCode(affCode string) (int, error) {
  206. if affCode == "" {
  207. return 0, errors.New("affCode 为空!")
  208. }
  209. var user User
  210. err := DB.Select("id").First(&user, "aff_code = ?", affCode).Error
  211. return user.Id, err
  212. }
  213. func DeleteUserById(id int) (err error) {
  214. if id == 0 {
  215. return errors.New("id 为空!")
  216. }
  217. user := User{Id: id}
  218. return user.Delete()
  219. }
  220. func HardDeleteUserById(id int) error {
  221. if id == 0 {
  222. return errors.New("id 为空!")
  223. }
  224. err := DB.Unscoped().Delete(&User{}, "id = ?", id).Error
  225. return err
  226. }
  227. func inviteUser(inviterId int) (err error) {
  228. user, err := GetUserById(inviterId, true)
  229. if err != nil {
  230. return err
  231. }
  232. user.AffCount++
  233. user.AffQuota += common.QuotaForInviter
  234. user.AffHistoryQuota += common.QuotaForInviter
  235. return DB.Save(user).Error
  236. }
  237. func (user *User) TransferAffQuotaToQuota(quota int) error {
  238. // 检查quota是否小于最小额度
  239. if float64(quota) < common.QuotaPerUnit {
  240. return fmt.Errorf("转移额度最小为%s!", common.LogQuota(int(common.QuotaPerUnit)))
  241. }
  242. // 开始数据库事务
  243. tx := DB.Begin()
  244. if tx.Error != nil {
  245. return tx.Error
  246. }
  247. defer tx.Rollback() // 确保在函数退出时事务能回滚
  248. // 加锁查询用户以确保数据一致性
  249. err := tx.Set("gorm:query_option", "FOR UPDATE").First(&user, user.Id).Error
  250. if err != nil {
  251. return err
  252. }
  253. // 再次检查用户的AffQuota是否足够
  254. if user.AffQuota < quota {
  255. return errors.New("邀请额度不足!")
  256. }
  257. // 更新用户额度
  258. user.AffQuota -= quota
  259. user.Quota += quota
  260. // 保存用户状态
  261. if err := tx.Save(user).Error; err != nil {
  262. return err
  263. }
  264. // 提交事务
  265. return tx.Commit().Error
  266. }
  267. func (user *User) Insert(inviterId int) error {
  268. var err error
  269. if user.Password != "" {
  270. user.Password, err = common.Password2Hash(user.Password)
  271. if err != nil {
  272. return err
  273. }
  274. }
  275. user.Quota = common.QuotaForNewUser
  276. //user.SetAccessToken(common.GetUUID())
  277. user.AffCode = common.GetRandomString(4)
  278. result := DB.Create(user)
  279. if result.Error != nil {
  280. return result.Error
  281. }
  282. if common.QuotaForNewUser > 0 {
  283. RecordLog(user.Id, LogTypeSystem, fmt.Sprintf("新用户注册赠送 %s", common.LogQuota(common.QuotaForNewUser)))
  284. }
  285. if inviterId != 0 {
  286. if common.QuotaForInvitee > 0 {
  287. _ = IncreaseUserQuota(user.Id, common.QuotaForInvitee, true)
  288. RecordLog(user.Id, LogTypeSystem, fmt.Sprintf("使用邀请码赠送 %s", common.LogQuota(common.QuotaForInvitee)))
  289. }
  290. if common.QuotaForInviter > 0 {
  291. //_ = IncreaseUserQuota(inviterId, common.QuotaForInviter)
  292. RecordLog(inviterId, LogTypeSystem, fmt.Sprintf("邀请用户赠送 %s", common.LogQuota(common.QuotaForInviter)))
  293. _ = inviteUser(inviterId)
  294. }
  295. }
  296. return nil
  297. }
  298. func (user *User) Update(updatePassword bool) error {
  299. var err error
  300. if updatePassword {
  301. user.Password, err = common.Password2Hash(user.Password)
  302. if err != nil {
  303. return err
  304. }
  305. }
  306. newUser := *user
  307. DB.First(&user, user.Id)
  308. if err = DB.Model(user).Updates(newUser).Error; err != nil {
  309. return err
  310. }
  311. // Update cache
  312. return updateUserCache(*user)
  313. }
  314. func (user *User) Edit(updatePassword bool) error {
  315. var err error
  316. if updatePassword {
  317. user.Password, err = common.Password2Hash(user.Password)
  318. if err != nil {
  319. return err
  320. }
  321. }
  322. newUser := *user
  323. updates := map[string]interface{}{
  324. "username": newUser.Username,
  325. "display_name": newUser.DisplayName,
  326. "group": newUser.Group,
  327. "quota": newUser.Quota,
  328. "remark": newUser.Remark,
  329. }
  330. if updatePassword {
  331. updates["password"] = newUser.Password
  332. }
  333. DB.First(&user, user.Id)
  334. if err = DB.Model(user).Updates(updates).Error; err != nil {
  335. return err
  336. }
  337. // Update cache
  338. return updateUserCache(*user)
  339. }
  340. func (user *User) Delete() error {
  341. if user.Id == 0 {
  342. return errors.New("id 为空!")
  343. }
  344. if err := DB.Delete(user).Error; err != nil {
  345. return err
  346. }
  347. // 清除缓存
  348. return invalidateUserCache(user.Id)
  349. }
  350. func (user *User) HardDelete() error {
  351. if user.Id == 0 {
  352. return errors.New("id 为空!")
  353. }
  354. err := DB.Unscoped().Delete(user).Error
  355. return err
  356. }
  357. // ValidateAndFill check password & user status
  358. func (user *User) ValidateAndFill() (err error) {
  359. // When querying with struct, GORM will only query with non-zero fields,
  360. // that means if your field's value is 0, '', false or other zero values,
  361. // it won't be used to build query conditions
  362. password := user.Password
  363. username := strings.TrimSpace(user.Username)
  364. if username == "" || password == "" {
  365. return errors.New("用户名或密码为空")
  366. }
  367. // find buy username or email
  368. DB.Where("username = ? OR email = ?", username, username).First(user)
  369. okay := common.ValidatePasswordAndHash(password, user.Password)
  370. if !okay || user.Status != common.UserStatusEnabled {
  371. return errors.New("用户名或密码错误,或用户已被封禁")
  372. }
  373. return nil
  374. }
  375. func (user *User) FillUserById() error {
  376. if user.Id == 0 {
  377. return errors.New("id 为空!")
  378. }
  379. DB.Where(User{Id: user.Id}).First(user)
  380. return nil
  381. }
  382. func (user *User) FillUserByEmail() error {
  383. if user.Email == "" {
  384. return errors.New("email 为空!")
  385. }
  386. DB.Where(User{Email: user.Email}).First(user)
  387. return nil
  388. }
  389. func (user *User) FillUserByGitHubId() error {
  390. if user.GitHubId == "" {
  391. return errors.New("GitHub id 为空!")
  392. }
  393. DB.Where(User{GitHubId: user.GitHubId}).First(user)
  394. return nil
  395. }
  396. func (user *User) FillUserByOidcId() error {
  397. if user.OidcId == "" {
  398. return errors.New("oidc id 为空!")
  399. }
  400. DB.Where(User{OidcId: user.OidcId}).First(user)
  401. return nil
  402. }
  403. func (user *User) FillUserByWeChatId() error {
  404. if user.WeChatId == "" {
  405. return errors.New("WeChat id 为空!")
  406. }
  407. DB.Where(User{WeChatId: user.WeChatId}).First(user)
  408. return nil
  409. }
  410. func (user *User) FillUserByTelegramId() error {
  411. if user.TelegramId == "" {
  412. return errors.New("Telegram id 为空!")
  413. }
  414. err := DB.Where(User{TelegramId: user.TelegramId}).First(user).Error
  415. if errors.Is(err, gorm.ErrRecordNotFound) {
  416. return errors.New("该 Telegram 账户未绑定")
  417. }
  418. return nil
  419. }
  420. func IsEmailAlreadyTaken(email string) bool {
  421. return DB.Unscoped().Where("email = ?", email).Find(&User{}).RowsAffected == 1
  422. }
  423. func IsWeChatIdAlreadyTaken(wechatId string) bool {
  424. return DB.Unscoped().Where("wechat_id = ?", wechatId).Find(&User{}).RowsAffected == 1
  425. }
  426. func IsGitHubIdAlreadyTaken(githubId string) bool {
  427. return DB.Unscoped().Where("github_id = ?", githubId).Find(&User{}).RowsAffected == 1
  428. }
  429. func IsOidcIdAlreadyTaken(oidcId string) bool {
  430. return DB.Where("oidc_id = ?", oidcId).Find(&User{}).RowsAffected == 1
  431. }
  432. func IsTelegramIdAlreadyTaken(telegramId string) bool {
  433. return DB.Unscoped().Where("telegram_id = ?", telegramId).Find(&User{}).RowsAffected == 1
  434. }
  435. func ResetUserPasswordByEmail(email string, password string) error {
  436. if email == "" || password == "" {
  437. return errors.New("邮箱地址或密码为空!")
  438. }
  439. hashedPassword, err := common.Password2Hash(password)
  440. if err != nil {
  441. return err
  442. }
  443. err = DB.Model(&User{}).Where("email = ?", email).Update("password", hashedPassword).Error
  444. return err
  445. }
  446. func IsAdmin(userId int) bool {
  447. if userId == 0 {
  448. return false
  449. }
  450. var user User
  451. err := DB.Where("id = ?", userId).Select("role").Find(&user).Error
  452. if err != nil {
  453. common.SysError("no such user " + err.Error())
  454. return false
  455. }
  456. return user.Role >= common.RoleAdminUser
  457. }
  458. //// IsUserEnabled checks user status from Redis first, falls back to DB if needed
  459. //func IsUserEnabled(id int, fromDB bool) (status bool, err error) {
  460. // defer func() {
  461. // // Update Redis cache asynchronously on successful DB read
  462. // if shouldUpdateRedis(fromDB, err) {
  463. // gopool.Go(func() {
  464. // if err := updateUserStatusCache(id, status); err != nil {
  465. // common.SysError("failed to update user status cache: " + err.Error())
  466. // }
  467. // })
  468. // }
  469. // }()
  470. // if !fromDB && common.RedisEnabled {
  471. // // Try Redis first
  472. // status, err := getUserStatusCache(id)
  473. // if err == nil {
  474. // return status == common.UserStatusEnabled, nil
  475. // }
  476. // // Don't return error - fall through to DB
  477. // }
  478. // fromDB = true
  479. // var user User
  480. // err = DB.Where("id = ?", id).Select("status").Find(&user).Error
  481. // if err != nil {
  482. // return false, err
  483. // }
  484. //
  485. // return user.Status == common.UserStatusEnabled, nil
  486. //}
  487. func ValidateAccessToken(token string) (user *User) {
  488. if token == "" {
  489. return nil
  490. }
  491. token = strings.Replace(token, "Bearer ", "", 1)
  492. user = &User{}
  493. if DB.Where("access_token = ?", token).First(user).RowsAffected == 1 {
  494. return user
  495. }
  496. return nil
  497. }
  498. // GetUserQuota gets quota from Redis first, falls back to DB if needed
  499. func GetUserQuota(id int, fromDB bool) (quota int, err error) {
  500. defer func() {
  501. // Update Redis cache asynchronously on successful DB read
  502. if shouldUpdateRedis(fromDB, err) {
  503. gopool.Go(func() {
  504. if err := updateUserQuotaCache(id, quota); err != nil {
  505. common.SysError("failed to update user quota cache: " + err.Error())
  506. }
  507. })
  508. }
  509. }()
  510. if !fromDB && common.RedisEnabled {
  511. quota, err := getUserQuotaCache(id)
  512. if err == nil {
  513. return quota, nil
  514. }
  515. // Don't return error - fall through to DB
  516. }
  517. fromDB = true
  518. err = DB.Model(&User{}).Where("id = ?", id).Select("quota").Find(&quota).Error
  519. if err != nil {
  520. return 0, err
  521. }
  522. return quota, nil
  523. }
  524. func GetUserUsedQuota(id int) (quota int, err error) {
  525. err = DB.Model(&User{}).Where("id = ?", id).Select("used_quota").Find(&quota).Error
  526. return quota, err
  527. }
  528. func GetUserEmail(id int) (email string, err error) {
  529. err = DB.Model(&User{}).Where("id = ?", id).Select("email").Find(&email).Error
  530. return email, err
  531. }
  532. // GetUserGroup gets group from Redis first, falls back to DB if needed
  533. func GetUserGroup(id int, fromDB bool) (group string, err error) {
  534. defer func() {
  535. // Update Redis cache asynchronously on successful DB read
  536. if shouldUpdateRedis(fromDB, err) {
  537. gopool.Go(func() {
  538. if err := updateUserGroupCache(id, group); err != nil {
  539. common.SysError("failed to update user group cache: " + err.Error())
  540. }
  541. })
  542. }
  543. }()
  544. if !fromDB && common.RedisEnabled {
  545. group, err := getUserGroupCache(id)
  546. if err == nil {
  547. return group, nil
  548. }
  549. // Don't return error - fall through to DB
  550. }
  551. fromDB = true
  552. err = DB.Model(&User{}).Where("id = ?", id).Select(commonGroupCol).Find(&group).Error
  553. if err != nil {
  554. return "", err
  555. }
  556. return group, nil
  557. }
  558. // GetUserSetting gets setting from Redis first, falls back to DB if needed
  559. func GetUserSetting(id int, fromDB bool) (settingMap map[string]interface{}, err error) {
  560. var setting string
  561. defer func() {
  562. // Update Redis cache asynchronously on successful DB read
  563. if shouldUpdateRedis(fromDB, err) {
  564. gopool.Go(func() {
  565. if err := updateUserSettingCache(id, setting); err != nil {
  566. common.SysError("failed to update user setting cache: " + err.Error())
  567. }
  568. })
  569. }
  570. }()
  571. if !fromDB && common.RedisEnabled {
  572. setting, err := getUserSettingCache(id)
  573. if err == nil {
  574. return setting, nil
  575. }
  576. // Don't return error - fall through to DB
  577. }
  578. fromDB = true
  579. err = DB.Model(&User{}).Where("id = ?", id).Select("setting").Find(&setting).Error
  580. if err != nil {
  581. return map[string]interface{}{}, err
  582. }
  583. return common.StrToMap(setting), nil
  584. }
  585. func IncreaseUserQuota(id int, quota int, db bool) (err error) {
  586. if quota < 0 {
  587. return errors.New("quota 不能为负数!")
  588. }
  589. gopool.Go(func() {
  590. err := cacheIncrUserQuota(id, int64(quota))
  591. if err != nil {
  592. common.SysError("failed to increase user quota: " + err.Error())
  593. }
  594. })
  595. if !db && common.BatchUpdateEnabled {
  596. addNewRecord(BatchUpdateTypeUserQuota, id, quota)
  597. return nil
  598. }
  599. return increaseUserQuota(id, quota)
  600. }
  601. func increaseUserQuota(id int, quota int) (err error) {
  602. err = DB.Model(&User{}).Where("id = ?", id).Update("quota", gorm.Expr("quota + ?", quota)).Error
  603. if err != nil {
  604. return err
  605. }
  606. return err
  607. }
  608. func DecreaseUserQuota(id int, quota int) (err error) {
  609. if quota < 0 {
  610. return errors.New("quota 不能为负数!")
  611. }
  612. gopool.Go(func() {
  613. err := cacheDecrUserQuota(id, int64(quota))
  614. if err != nil {
  615. common.SysError("failed to decrease user quota: " + err.Error())
  616. }
  617. })
  618. if common.BatchUpdateEnabled {
  619. addNewRecord(BatchUpdateTypeUserQuota, id, -quota)
  620. return nil
  621. }
  622. return decreaseUserQuota(id, quota)
  623. }
  624. func decreaseUserQuota(id int, quota int) (err error) {
  625. err = DB.Model(&User{}).Where("id = ?", id).Update("quota", gorm.Expr("quota - ?", quota)).Error
  626. if err != nil {
  627. return err
  628. }
  629. return err
  630. }
  631. func DeltaUpdateUserQuota(id int, delta int) (err error) {
  632. if delta == 0 {
  633. return nil
  634. }
  635. if delta > 0 {
  636. return IncreaseUserQuota(id, delta, false)
  637. } else {
  638. return DecreaseUserQuota(id, -delta)
  639. }
  640. }
  641. //func GetRootUserEmail() (email string) {
  642. // DB.Model(&User{}).Where("role = ?", common.RoleRootUser).Select("email").Find(&email)
  643. // return email
  644. //}
  645. func GetRootUser() (user *User) {
  646. DB.Where("role = ?", common.RoleRootUser).First(&user)
  647. return user
  648. }
  649. func UpdateUserUsedQuotaAndRequestCount(id int, quota int) {
  650. if common.BatchUpdateEnabled {
  651. addNewRecord(BatchUpdateTypeUsedQuota, id, quota)
  652. addNewRecord(BatchUpdateTypeRequestCount, id, 1)
  653. return
  654. }
  655. updateUserUsedQuotaAndRequestCount(id, quota, 1)
  656. }
  657. func updateUserUsedQuotaAndRequestCount(id int, quota int, count int) {
  658. err := DB.Model(&User{}).Where("id = ?", id).Updates(
  659. map[string]interface{}{
  660. "used_quota": gorm.Expr("used_quota + ?", quota),
  661. "request_count": gorm.Expr("request_count + ?", count),
  662. },
  663. ).Error
  664. if err != nil {
  665. common.SysError("failed to update user used quota and request count: " + err.Error())
  666. return
  667. }
  668. //// 更新缓存
  669. //if err := invalidateUserCache(id); err != nil {
  670. // common.SysError("failed to invalidate user cache: " + err.Error())
  671. //}
  672. }
  673. func updateUserUsedQuota(id int, quota int) {
  674. err := DB.Model(&User{}).Where("id = ?", id).Updates(
  675. map[string]interface{}{
  676. "used_quota": gorm.Expr("used_quota + ?", quota),
  677. },
  678. ).Error
  679. if err != nil {
  680. common.SysError("failed to update user used quota: " + err.Error())
  681. }
  682. }
  683. func updateUserRequestCount(id int, count int) {
  684. err := DB.Model(&User{}).Where("id = ?", id).Update("request_count", gorm.Expr("request_count + ?", count)).Error
  685. if err != nil {
  686. common.SysError("failed to update user request count: " + err.Error())
  687. }
  688. }
  689. // GetUsernameById gets username from Redis first, falls back to DB if needed
  690. func GetUsernameById(id int, fromDB bool) (username string, err error) {
  691. defer func() {
  692. // Update Redis cache asynchronously on successful DB read
  693. if shouldUpdateRedis(fromDB, err) {
  694. gopool.Go(func() {
  695. if err := updateUserNameCache(id, username); err != nil {
  696. common.SysError("failed to update user name cache: " + err.Error())
  697. }
  698. })
  699. }
  700. }()
  701. if !fromDB && common.RedisEnabled {
  702. username, err := getUserNameCache(id)
  703. if err == nil {
  704. return username, nil
  705. }
  706. // Don't return error - fall through to DB
  707. }
  708. fromDB = true
  709. err = DB.Model(&User{}).Where("id = ?", id).Select("username").Find(&username).Error
  710. if err != nil {
  711. return "", err
  712. }
  713. return username, nil
  714. }
  715. func IsLinuxDOIdAlreadyTaken(linuxDOId string) bool {
  716. var user User
  717. err := DB.Unscoped().Where("linux_do_id = ?", linuxDOId).First(&user).Error
  718. return !errors.Is(err, gorm.ErrRecordNotFound)
  719. }
  720. func (user *User) FillUserByLinuxDOId() error {
  721. if user.LinuxDOId == "" {
  722. return errors.New("linux do id is empty")
  723. }
  724. err := DB.Where("linux_do_id = ?", user.LinuxDOId).First(user).Error
  725. return err
  726. }
  727. func RootUserExists() bool {
  728. var user User
  729. err := DB.Where("role = ?", common.RoleRootUser).First(&user).Error
  730. if err != nil {
  731. return false
  732. }
  733. return true
  734. }