user.go 23 KB

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