gemini_token_cache.go 1.4 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849
  1. package repository
  2. import (
  3. "context"
  4. "fmt"
  5. "time"
  6. "github.com/Wei-Shaw/sub2api/internal/service"
  7. "github.com/redis/go-redis/v9"
  8. )
  9. const (
  10. oauthTokenKeyPrefix = "oauth:token:"
  11. oauthRefreshLockKeyPrefix = "oauth:refresh_lock:"
  12. )
  13. type geminiTokenCache struct {
  14. rdb *redis.Client
  15. }
  16. func NewGeminiTokenCache(rdb *redis.Client) service.GeminiTokenCache {
  17. return &geminiTokenCache{rdb: rdb}
  18. }
  19. func (c *geminiTokenCache) GetAccessToken(ctx context.Context, cacheKey string) (string, error) {
  20. key := fmt.Sprintf("%s%s", oauthTokenKeyPrefix, cacheKey)
  21. return c.rdb.Get(ctx, key).Result()
  22. }
  23. func (c *geminiTokenCache) SetAccessToken(ctx context.Context, cacheKey string, token string, ttl time.Duration) error {
  24. key := fmt.Sprintf("%s%s", oauthTokenKeyPrefix, cacheKey)
  25. return c.rdb.Set(ctx, key, token, ttl).Err()
  26. }
  27. func (c *geminiTokenCache) DeleteAccessToken(ctx context.Context, cacheKey string) error {
  28. key := fmt.Sprintf("%s%s", oauthTokenKeyPrefix, cacheKey)
  29. return c.rdb.Del(ctx, key).Err()
  30. }
  31. func (c *geminiTokenCache) AcquireRefreshLock(ctx context.Context, cacheKey string, ttl time.Duration) (bool, error) {
  32. key := fmt.Sprintf("%s%s", oauthRefreshLockKeyPrefix, cacheKey)
  33. return c.rdb.SetNX(ctx, key, 1, ttl).Result()
  34. }
  35. func (c *geminiTokenCache) ReleaseRefreshLock(ctx context.Context, cacheKey string) error {
  36. key := fmt.Sprintf("%s%s", oauthRefreshLockKeyPrefix, cacheKey)
  37. return c.rdb.Del(ctx, key).Err()
  38. }