2025-05-07 14:12:53 +08:00

143 lines
3.6 KiB
Go

package redis
import (
"bytes"
"context"
"encoding/gob"
"errors"
"fmt"
"time"
"management/internal/pkg/config"
"github.com/drhin/logger"
"github.com/redis/go-redis/v9"
)
var ErrRedisKeyNotFound = errors.New("redis key not found")
type Cache interface {
Encode(a any) ([]byte, error)
Set(ctx context.Context, key string, value interface{}, expiration time.Duration) error
Del(ctx context.Context, keys ...string) error
Get(ctx context.Context, key string) (string, error)
GetBytes(ctx context.Context, key string) ([]byte, error)
Scan(ctx context.Context, cursor uint64, match string, count int64) *redis.ScanCmd
Keys(ctx context.Context, pattern string) ([]string, error)
ListKeys(ctx context.Context, pattern string, pageID int, pageSize int) ([]string, int, error)
}
type redisCache struct {
client *redis.Client
}
func New(conf *config.Config, log *logger.Logger) (Cache, func(), error) {
rdb := redis.NewClient(&redis.Options{
Addr: fmt.Sprintf("%s:%d", conf.Redis.Host, conf.Redis.Port),
Password: conf.Redis.Password,
DB: conf.Redis.DB,
})
_, err := rdb.Ping(context.Background()).Result()
if err != nil {
return nil, nil, err
}
cleanup := func() {
if err := rdb.Close(); err != nil {
log.Error("redis close error", err)
}
}
return &redisCache{
client: rdb,
}, cleanup, nil
}
func (r *redisCache) Encode(a any) ([]byte, error) {
var b bytes.Buffer
if err := gob.NewEncoder(&b).Encode(a); err != nil {
return nil, err
}
return b.Bytes(), nil
}
// Set 设置值
func (r *redisCache) Set(ctx context.Context, key string, value interface{}, expiration time.Duration) error {
return r.client.Set(ctx, key, value, expiration).Err()
}
// Del 删除键值
func (r *redisCache) Del(ctx context.Context, keys ...string) error {
return r.client.Del(ctx, keys...).Err()
}
// Get 获取值
func (r *redisCache) Get(ctx context.Context, key string) (string, error) {
val, err := r.client.Get(ctx, key).Result()
if errors.Is(err, redis.Nil) {
return "", ErrRedisKeyNotFound
} else if err != nil {
return "", fmt.Errorf("cannot get value with:[%s]: %v", key, err)
} else {
return val, nil
}
}
// GetBytes 获取值
func (r *redisCache) GetBytes(ctx context.Context, key string) ([]byte, error) {
val, err := r.client.Get(ctx, key).Bytes()
if errors.Is(err, redis.Nil) {
return nil, ErrRedisKeyNotFound
} else if err != nil {
return nil, fmt.Errorf("cannot get value with:[%s]: %v", key, err)
} else {
return val, nil
}
}
func (r *redisCache) Scan(ctx context.Context, cursor uint64, match string, count int64) *redis.ScanCmd {
return r.client.Scan(ctx, cursor, match, count)
}
func (r *redisCache) Keys(ctx context.Context, pattern string) ([]string, error) {
return r.client.Keys(ctx, pattern).Result()
}
func (r *redisCache) ListKeys(ctx context.Context, pattern string, pageID int, pageSize int) ([]string, int, error) {
all, err := r.client.Keys(ctx, pattern).Result()
if err != nil {
return nil, 0, err
}
count := len(all)
if count == 0 {
return nil, 0, err
}
// 使用SCAN命令分页获取键
cursor := uint64(0)
var keys []string
for {
var scanResult []string
var err error
scanResult, cursor, err = r.client.Scan(ctx, cursor, pattern, int64(pageSize)).Result()
if err != nil {
return nil, count, err
}
keys = append(keys, scanResult...)
if cursor == 0 {
break
}
}
startIndex := (pageID - 1) * pageSize
endIndex := startIndex + pageSize
if startIndex >= len(keys) {
return nil, count, nil
}
if endIndex > len(keys) {
endIndex = len(keys)
}
return keys[startIndex:endIndex], count, nil
}