diff --git a/abstract/IRedisClient.go b/abstract/IRedisClient.go index 55eafad..b10f2b3 100644 --- a/abstract/IRedisClient.go +++ b/abstract/IRedisClient.go @@ -8,6 +8,8 @@ package abstract import ( + "context" + "git.zhangdeman.cn/zhangdeman/redis/define" "github.com/redis/go-redis/v9" "go.uber.org/zap" ) @@ -18,9 +20,11 @@ import ( // // Date : 15:09 2024/6/18 type IRedisClient interface { - GetRealClient(instanceFlag string) *redis.ClusterClient // 获取客户端连接 - GetRealClientWithError(instanceFlag string) (*redis.ClusterClient, error) // 获取带error的客户端连接 - AddClient(instanceFlag string, instanceConfig *redis.ClusterOptions) error // 添加新的客户端连接 - RemoveClient(instanceFlag string) // 移除一个客户端连接 - SetLogger(loggerInstance *zap.Logger) // 设置日志实例 + GetRealClient(instanceFlag string) *redis.ClusterClient // 获取客户端连接 + GetRealClientWithError(instanceFlag string) (*redis.ClusterClient, error) // 获取带error的客户端连接 + AddClient(instanceFlag string, instanceConfig *redis.ClusterOptions) error // 添加新的客户端连接 + RemoveClient(instanceFlag string) // 移除一个客户端连接 + SetLogger(loggerInstance *zap.Logger) // 设置日志实例, 全局生效, 而非针对某一个实例 + Exec(ctx context.Context, instanceFlag string, command string, args ...any) *define.RedisResult // 执行任意命令 + SetCommandWhiteList(command []string) // 设置命令的白名单,全局生效, 而非单独针对某一个实例 } diff --git a/client.go b/client.go index 712bf27..de19571 100644 --- a/client.go +++ b/client.go @@ -8,26 +8,120 @@ package redis import ( + "context" "git.zhangdeman.cn/zhangdeman/redis/abstract" + "git.zhangdeman.cn/zhangdeman/redis/define" + "git.zhangdeman.cn/zhangdeman/wrapper" "github.com/pkg/errors" redisClient "github.com/redis/go-redis/v9" "go.uber.org/zap" + "strings" "sync" + "time" ) var Client abstract.IRedisClient func init() { Client = &OwnClient{ - lock: &sync.RWMutex{}, - instanceTable: make(map[string]*redisClient.ClusterClient), + lock: &sync.RWMutex{}, + instanceTable: make(map[string]*redisClient.ClusterClient), + whiteCommandTable: make(map[string]bool), } } type OwnClient struct { - lock *sync.RWMutex - instanceTable map[string]*redisClient.ClusterClient - logger *zap.Logger + lock *sync.RWMutex + instanceTable map[string]*redisClient.ClusterClient + whiteCommandTable map[string]bool + logger *zap.Logger +} + +func (o *OwnClient) isAllowCommand(command string) bool { + if len(o.whiteCommandTable) == 0 { + // 未配置, 视为全部允许执行 + return true + } + if o.whiteCommandTable["*"] { + // 配置了 * 视为全部允许执行 + return true + } + command = strings.ToLower(strings.TrimSpace(command)) + o.lock.RLock() + defer o.lock.RUnlock() + return o.whiteCommandTable[command] +} + +// Exec 执行命令 +// +// Author : go_developer@163.com<白茶清欢> +// +// Date : 11:05 2024/6/19 +func (o *OwnClient) Exec(ctx context.Context, instanceFlag string, command string, args ...any) *define.RedisResult { + var ( + instance *redisClient.ClusterClient + ) + + cmdParamList := []any{ + command, + } + argStrList := make([]string, 0) + for _, itemArg := range args { + argStrList = append(argStrList, wrapper.AnyDataType(itemArg).ToString().Value()) + cmdParamList = append(cmdParamList, itemArg) + } + res := &define.RedisResult{ + StartTime: time.Now().UnixMilli(), + FinishTime: 0, + UsedTime: 0, + Result: "", + Command: command, + ArgList: argStrList, + Err: nil, + } + defer func() { + res.FinishTime = time.Now().UnixMilli() + res.UsedTime = res.FinishTime - res.StartTime + if nil == o.logger { + // 未注入日志实例 + return + } + o.logger.Info( + "Redis命令执行记录", + zap.Int64("start_time", res.StartTime), + zap.Int64("finish_time", res.FinishTime), + zap.Int64("used_time", res.UsedTime), + zap.String("command", res.Command), + zap.String("arg_list", strings.Join(res.ArgList, " ")), + zap.String("execute_result", res.Result), + zap.Error(res.Err), + ) + }() + if instance, res.Err = o.GetRealClientWithError(instanceFlag); nil != res.Err { + return res + } + if nil == ctx { + ctx = context.Background() + } + cmdRes := instance.Do(ctx, cmdParamList...) + if res.Err = cmdRes.Err(); nil != res.Err { + return res + } + res.Result = cmdRes.String() + return res +} + +// SetCommandWhiteList 设置命令白名单, 空 或者 包含 * 则认为所有命令均允许执行 +// +// Author : go_developer@163.com<白茶清欢> +// +// Date : 11:02 2024/6/19 +func (o *OwnClient) SetCommandWhiteList(commandList []string) { + o.lock.Lock() + defer o.lock.Unlock() + for _, itemCommand := range commandList { + o.whiteCommandTable[strings.ToLower(strings.TrimSpace(itemCommand))] = true + } } func (o *OwnClient) GetRealClient(instanceFlag string) *redisClient.ClusterClient { diff --git a/define/result.go b/define/result.go new file mode 100644 index 0000000..e4599b7 --- /dev/null +++ b/define/result.go @@ -0,0 +1,24 @@ +// Package define ... +// +// Description : define ... +// +// Author : go_developer@163.com<白茶清欢> +// +// Date : 2024-06-19 10:49 +package define + +// RedisResult redis名玲玲执行结果 +// +// Author : go_developer@163.com<白茶清欢> +// +// Date : 10:50 2024/6/19 +type RedisResult struct { + InstanceFlag string `json:"instance_flag"` // 实例标识 + StartTime int64 `json:"start_time"` // 开始执行时间, 单位 : ms + FinishTime int64 `json:"finish_time"` // 完成执行时间, 单位 : ms + UsedTime int64 `json:"used_time"` // 执行耗时, 单位 : ms + Result string `json:"result"` // 执行结果 + Command string `json:"command"` // 执行的命令 + ArgList []string `json:"arg_list"` // 参数列表 + Err error `json:"-"` // 失败信息 +}