user.go 23 KB

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