kra/internal/initialize/redis.go

137 lines
2.6 KiB
Go

package initialize
import (
"context"
"fmt"
"kra/internal/conf"
"github.com/redis/go-redis/v9"
)
// RedisClient Redis客户端包装器
type RedisClient struct {
client redis.UniversalClient
name string
}
// NewRedisClient 创建Redis客户端
func NewRedisClient(cfg *conf.Redis) (*RedisClient, error) {
if cfg == nil {
return nil, fmt.Errorf("redis config is nil")
}
client, err := initRedisClient(cfg)
if err != nil {
return nil, err
}
return &RedisClient{
client: client,
name: cfg.Name,
}, nil
}
// initRedisClient 初始化Redis客户端
func initRedisClient(cfg *conf.Redis) (redis.UniversalClient, error) {
var client redis.UniversalClient
// 使用集群模式
if cfg.UseCluster {
client = redis.NewClusterClient(&redis.ClusterOptions{
Addrs: cfg.ClusterAddrs,
Password: cfg.Password,
})
} else {
// 使用单例模式
client = redis.NewClient(&redis.Options{
Addr: cfg.Addr,
Password: cfg.Password,
DB: int(cfg.Db),
})
}
// 测试连接
pong, err := client.Ping(context.Background()).Result()
if err != nil {
return nil, fmt.Errorf("redis connect ping failed, name: %s, err: %v", cfg.Name, err)
}
_ = pong // 连接成功
return client, nil
}
// Client 获取Redis客户端
func (r *RedisClient) Client() redis.UniversalClient {
return r.client
}
// Name 获取Redis名称
func (r *RedisClient) Name() string {
return r.name
}
// Close 关闭Redis连接
func (r *RedisClient) Close() error {
if r.client != nil {
return r.client.Close()
}
return nil
}
// RedisList Redis列表管理器
type RedisList struct {
clients map[string]*RedisClient
}
// NewRedisList 创建Redis列表
func NewRedisList() *RedisList {
return &RedisList{
clients: make(map[string]*RedisClient),
}
}
// InitRedisList 初始化多Redis连接
func InitRedisList(redisList []*conf.Redis) (*RedisList, error) {
list := NewRedisList()
for _, cfg := range redisList {
client, err := NewRedisClient(cfg)
if err != nil {
return nil, err
}
list.clients[cfg.Name] = client
}
return list, nil
}
// Get 获取指定名称的Redis客户端
func (l *RedisList) Get(name string) *RedisClient {
if l.clients == nil {
return nil
}
return l.clients[name]
}
// Set 设置Redis客户端
func (l *RedisList) Set(name string, client *RedisClient) {
if l.clients == nil {
l.clients = make(map[string]*RedisClient)
}
l.clients[name] = client
}
// All 获取所有Redis客户端
func (l *RedisList) All() map[string]*RedisClient {
return l.clients
}
// Close 关闭所有Redis连接
func (l *RedisList) Close() error {
for _, client := range l.clients {
client.Close()
}
return nil
}