You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
 
 
 

237 line
5.9 KiB

  1. package model
  2. import (
  3. "fmt"
  4. "time"
  5. "github.com/QuantumNous/new-api/common"
  6. "github.com/QuantumNous/new-api/constant"
  7. "github.com/QuantumNous/new-api/dto"
  8. "github.com/gin-gonic/gin"
  9. "github.com/bytedance/gopkg/util/gopool"
  10. )
  11. // UserBase struct remains the same as it represents the cached data structure
  12. type UserBase struct {
  13. Source string `json:"source"`
  14. Id int `json:"id"`
  15. Group string `json:"group"`
  16. Email string `json:"email"`
  17. Quota int `json:"quota"`
  18. Status int `json:"status"`
  19. Username string `json:"username"`
  20. Setting string `json:"setting"`
  21. }
  22. func (user *UserBase) WriteContext(c *gin.Context) {
  23. common.SetContextKey(c, constant.ContextKeyUserGroup, user.Group)
  24. common.SetContextKey(c, constant.ContextKeyUserQuota, user.Quota)
  25. common.SetContextKey(c, constant.ContextKeyUserStatus, user.Status)
  26. common.SetContextKey(c, constant.ContextKeyUserEmail, user.Email)
  27. common.SetContextKey(c, constant.ContextKeyUserName, user.Username)
  28. common.SetContextKey(c, constant.ContextKeyUserSetting, user.GetSetting())
  29. common.SetContextKey(c, constant.ContextKeyUserSource, user.Source)
  30. }
  31. func (user *UserBase) GetSetting() dto.UserSetting {
  32. setting := dto.UserSetting{}
  33. if user.Setting != "" {
  34. err := common.Unmarshal([]byte(user.Setting), &setting)
  35. if err != nil {
  36. common.SysLog("failed to unmarshal setting: " + err.Error())
  37. }
  38. }
  39. return setting
  40. }
  41. // getUserCacheKey returns the key for user cache
  42. func getUserCacheKey(userId int) string {
  43. return fmt.Sprintf("user:%d", userId)
  44. }
  45. // invalidateUserCache clears user cache
  46. func invalidateUserCache(userId int) error {
  47. if !common.RedisEnabled {
  48. return nil
  49. }
  50. return common.RedisDelKey(getUserCacheKey(userId))
  51. }
  52. // updateUserCache updates all user cache fields using hash
  53. func updateUserCache(user User) error {
  54. if !common.RedisEnabled {
  55. return nil
  56. }
  57. return common.RedisHSetObj(
  58. getUserCacheKey(user.Id),
  59. user.ToBaseUser(),
  60. time.Duration(common.RedisKeyCacheSeconds())*time.Second,
  61. )
  62. }
  63. // GetUserCache gets complete user cache from hash
  64. func GetUserCache(userId int) (userCache *UserBase, err error) {
  65. var user *User
  66. var fromDB bool
  67. defer func() {
  68. // Update Redis cache asynchronously on successful DB read
  69. if shouldUpdateRedis(fromDB, err) && user != nil {
  70. gopool.Go(func() {
  71. if err := updateUserCache(*user); err != nil {
  72. common.SysLog("failed to update user status cache: " + err.Error())
  73. }
  74. })
  75. }
  76. }()
  77. // Try getting from Redis first
  78. userCache, err = cacheGetUserBase(userId)
  79. if err == nil {
  80. return userCache, nil
  81. }
  82. // If Redis fails, get from DB
  83. fromDB = true
  84. user, err = GetUserById(userId, false)
  85. if err != nil {
  86. return nil, err // Return nil and error if DB lookup fails
  87. }
  88. // Create cache object from user data
  89. userCache = &UserBase{
  90. Id: user.Id,
  91. Source: user.Source,
  92. Group: user.Group,
  93. Quota: user.Quota,
  94. Status: user.Status,
  95. Username: user.Username,
  96. Setting: user.Setting,
  97. Email: user.Email,
  98. }
  99. return userCache, nil
  100. }
  101. func cacheGetUserBase(userId int) (*UserBase, error) {
  102. if !common.RedisEnabled {
  103. return nil, fmt.Errorf("redis is not enabled")
  104. }
  105. var userCache UserBase
  106. // Try getting from Redis first
  107. err := common.RedisHGetObj(getUserCacheKey(userId), &userCache)
  108. if err != nil {
  109. return nil, err
  110. }
  111. return &userCache, nil
  112. }
  113. // Add atomic quota operations using hash fields
  114. func cacheIncrUserQuota(userId int, delta int64) error {
  115. if !common.RedisEnabled {
  116. return nil
  117. }
  118. return common.RedisHIncrBy(getUserCacheKey(userId), "Quota", delta)
  119. }
  120. func cacheDecrUserQuota(userId int, delta int64) error {
  121. return cacheIncrUserQuota(userId, -delta)
  122. }
  123. // Helper functions to get individual fields if needed
  124. func getUserGroupCache(userId int) (string, error) {
  125. cache, err := GetUserCache(userId)
  126. if err != nil {
  127. return "", err
  128. }
  129. return cache.Group, nil
  130. }
  131. func getUserQuotaCache(userId int) (int, error) {
  132. cache, err := GetUserCache(userId)
  133. if err != nil {
  134. return 0, err
  135. }
  136. return cache.Quota, nil
  137. }
  138. func getUserStatusCache(userId int) (int, error) {
  139. cache, err := GetUserCache(userId)
  140. if err != nil {
  141. return 0, err
  142. }
  143. return cache.Status, nil
  144. }
  145. func getUserNameCache(userId int) (string, error) {
  146. cache, err := GetUserCache(userId)
  147. if err != nil {
  148. return "", err
  149. }
  150. return cache.Username, nil
  151. }
  152. func getUserSettingCache(userId int) (dto.UserSetting, error) {
  153. cache, err := GetUserCache(userId)
  154. if err != nil {
  155. return dto.UserSetting{}, err
  156. }
  157. return cache.GetSetting(), nil
  158. }
  159. // New functions for individual field updates
  160. func updateUserStatusCache(userId int, status bool) error {
  161. if !common.RedisEnabled {
  162. return nil
  163. }
  164. statusInt := common.UserStatusEnabled
  165. if !status {
  166. statusInt = common.UserStatusDisabled
  167. }
  168. return common.RedisHSetField(getUserCacheKey(userId), "Status", fmt.Sprintf("%d", statusInt))
  169. }
  170. func updateUserQuotaCache(userId int, quota int) error {
  171. if !common.RedisEnabled {
  172. return nil
  173. }
  174. return common.RedisHSetField(getUserCacheKey(userId), "Quota", fmt.Sprintf("%d", quota))
  175. }
  176. func updateUserGroupCache(userId int, group string) error {
  177. if !common.RedisEnabled {
  178. return nil
  179. }
  180. return common.RedisHSetField(getUserCacheKey(userId), "Group", group)
  181. }
  182. func UpdateUserGroupCache(userId int, group string) error {
  183. return updateUserGroupCache(userId, group)
  184. }
  185. func updateUserNameCache(userId int, username string) error {
  186. if !common.RedisEnabled {
  187. return nil
  188. }
  189. return common.RedisHSetField(getUserCacheKey(userId), "Username", username)
  190. }
  191. func updateUserSettingCache(userId int, setting string) error {
  192. if !common.RedisEnabled {
  193. return nil
  194. }
  195. return common.RedisHSetField(getUserCacheKey(userId), "Setting", setting)
  196. }
  197. // GetUserLanguage returns the user's language preference from cache
  198. // Uses the existing GetUserCache mechanism for efficiency
  199. func GetUserLanguage(userId int) string {
  200. userCache, err := GetUserCache(userId)
  201. if err != nil {
  202. return ""
  203. }
  204. return userCache.GetSetting().Language
  205. }