user.go 22 KB

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