mirror of
https://github.com/charlienet/go-mixed.git
synced 2025-07-17 16:12:42 +08:00
98 lines
2.1 KiB
Go
98 lines
2.1 KiB
Go
package redis
|
||
|
||
import (
|
||
"context"
|
||
"net"
|
||
"strings"
|
||
|
||
"github.com/redis/go-redis/v9"
|
||
)
|
||
|
||
type renameHook struct {
|
||
prefix redisPrefix
|
||
}
|
||
|
||
func (r renameHook) DialHook(next redis.DialHook) redis.DialHook {
|
||
return func(ctx context.Context, network, addr string) (net.Conn, error) {
|
||
return next(ctx, network, addr)
|
||
}
|
||
}
|
||
|
||
func (r renameHook) ProcessPipelineHook(next redis.ProcessPipelineHook) redis.ProcessPipelineHook {
|
||
return func(ctx context.Context, cmds []redis.Cmder) error {
|
||
|
||
// 对多个KEY进行更名操作
|
||
for i := 0; i < len(cmds); i++ {
|
||
r.renameKey(cmds[i])
|
||
}
|
||
|
||
return next(ctx, cmds)
|
||
}
|
||
}
|
||
|
||
func (r renameHook) ProcessHook(next redis.ProcessHook) redis.ProcessHook {
|
||
return func(ctx context.Context, cmd redis.Cmder) error {
|
||
r.renameKey(cmd)
|
||
return next(ctx, cmd)
|
||
}
|
||
}
|
||
|
||
func (r renameHook) renameKey(cmd redis.Cmder) {
|
||
if !r.prefix.hasPrefix() {
|
||
return
|
||
}
|
||
|
||
args := cmd.Args()
|
||
if len(args) == 1 {
|
||
return
|
||
}
|
||
|
||
switch strings.ToUpper(cmd.Name()) {
|
||
case "SELECT", "FUNCTION", "INFO":
|
||
// 无KEY指令
|
||
case
|
||
"RENAME", "RENAMENX",
|
||
"MGET",
|
||
"RPOPLPUSH",
|
||
"SDIFF", "SDIFFSTORE", "SINTER", "SINTERSTORE",
|
||
"SUNION", "SUNIONSTORE",
|
||
"WATCH":
|
||
// 连续KEY
|
||
r.rename(args, createSepuence(1, len(args), 1)...)
|
||
case
|
||
"BLPOP", "BRPOP",
|
||
"BRPOPLPUSH ",
|
||
"SMOVE":
|
||
// 除最后一个外连续键
|
||
r.rename(args, createSepuence(1, len(args)-1, 1)...)
|
||
case "MSET", "MSETNX":
|
||
// 间隔KEY，KEY位置规则1,3,5,7
|
||
r.rename(args, createSepuence(1, len(args), 2)...)
|
||
case "EVAL", "EVALSHA", "EVALSHA_RO", "FCALL", "FCALL_RO":
|
||
// 命令中包含键数量
|
||
if n, ok := args[2].(int); ok && n > 0 {
|
||
r.rename(args, createSepuence(3, 3+n, 1)...)
|
||
}
|
||
default:
|
||
// 默认第一个参数为键值
|
||
r.rename(args, 1)
|
||
}
|
||
}
|
||
|
||
func (r renameHook) rename(args []any, indexes ...int) {
|
||
for _, i := range indexes {
|
||
if key, ok := args[i].(string); ok {
|
||
newKey := r.prefix.join(key)
|
||
args[i] = newKey
|
||
}
|
||
}
|
||
}
|
||
|
||
func createSepuence(start, end, step int) []int {
|
||
ret := make([]int, 0, (end-start)/step+1)
|
||
for i := start; i < end; i += step {
|
||
ret = append(ret, i)
|
||
}
|
||
return ret
|
||
}
|