user.go 23 KB

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