diff --git a/src/redis-shake/common/utils.go b/src/redis-shake/common/utils.go index 1d5712e..8314cd5 100644 --- a/src/redis-shake/common/utils.go +++ b/src/redis-shake/common/utils.go @@ -35,7 +35,7 @@ func OpenRedisConn(target []string, auth_type, passwd string, isCluster bool, tl } func OpenRedisConnWithTimeout(target []string, auth_type, passwd string, readTimeout, writeTimeout time.Duration, - isCluster bool, tlsEnable bool) redigo.Conn { + isCluster bool, tlsEnable bool) redigo.Conn { // return redigo.NewConn(OpenNetConn(target, auth_type, passwd), readTimeout, writeTimeout) if isCluster { cluster, err := redigoCluster.NewCluster( @@ -1003,3 +1003,13 @@ var defaultDialFunction = func(addr string) (redigo.Conn, error) { } return c, nil } + +// HasAtLeastOnePrefix checks whether the key begins with at least one of prefixes. +func HasAtLeastOnePrefix(key string, prefixes []string) bool { + for _, prefix := range prefixes { + if strings.HasPrefix(key, prefix) { + return true + } + } + return false +} diff --git a/src/redis-shake/common/utils_test.go b/src/redis-shake/common/utils_test.go index e1a571c..88986e5 100644 --- a/src/redis-shake/common/utils_test.go +++ b/src/redis-shake/common/utils_test.go @@ -1,9 +1,9 @@ package utils import ( - "testing" "fmt" "sort" + "testing" "github.com/stretchr/testify/assert" ) @@ -14,8 +14,8 @@ func TestGetAllClusterNode(t *testing.T) { fmt.Printf("TestGetAllClusterNode case %d.\n", nr) nr++ - client := OpenRedisConn([]string{"10.1.1.1:21333"}, "auth", "123456", false) - ret, err := GetAllClusterNode(client, "master") + client := OpenRedisConn([]string{"10.1.1.1:21333"}, "auth", "123456", false, false) + ret, err := GetAllClusterNode(client, "master", "") sort.Strings(ret) assert.Equal(t, nil, err, "should be equal") assert.Equal(t, 3, len(ret), "should be equal") @@ -28,8 +28,8 @@ func TestGetAllClusterNode(t *testing.T) { fmt.Printf("TestGetAllClusterNode case %d.\n", nr) nr++ - client := OpenRedisConn([]string{"10.1.1.1:21333"}, "auth", "123456", false) - ret, err := GetAllClusterNode(client, "slave") + client := OpenRedisConn([]string{"10.1.1.1:21333"}, "auth", "123456", false, false) + ret, err := GetAllClusterNode(client, "slave", "") sort.Strings(ret) assert.Equal(t, nil, err, "should be equal") assert.Equal(t, 3, len(ret), "should be equal") @@ -42,8 +42,8 @@ func TestGetAllClusterNode(t *testing.T) { fmt.Printf("TestGetAllClusterNode case %d.\n", nr) nr++ - client := OpenRedisConn([]string{"10.1.1.1:21333"}, "auth", "123456", false) - ret, err := GetAllClusterNode(client, "all") + client := OpenRedisConn([]string{"10.1.1.1:21333"}, "auth", "123456", false, false) + ret, err := GetAllClusterNode(client, "all", "") sort.Strings(ret) assert.Equal(t, nil, err, "should be equal") assert.Equal(t, 6, len(ret), "should be equal") @@ -54,4 +54,36 @@ func TestGetAllClusterNode(t *testing.T) { assert.Equal(t, "10.1.1.1:21335", ret[4], "should be equal") assert.Equal(t, "10.1.1.1:21336", ret[5], "should be equal") } -} \ No newline at end of file +} + +func TestHasAtLeastOnePrefix(t *testing.T) { + cases := []struct { + key string + prefixes []string + expectResult bool + }{ + { + // no prefix provided + "a", + []string{}, + false, + }, + { + // has prefix + "abc", + []string{"ab"}, + true, + }, + { + // does NOT have prefix + "abc", + []string{"edf", "wab"}, + false, + }, + } + + for _, c := range cases { + result := HasAtLeastOnePrefix(c.key, c.prefixes) + assert.Equal(t, c.expectResult, result) + } +} diff --git a/src/redis-shake/restore.go b/src/redis-shake/restore.go index 5eb39a7..3b63c3a 100644 --- a/src/redis-shake/restore.go +++ b/src/redis-shake/restore.go @@ -167,7 +167,7 @@ func (dr *dbRestorer) restoreRDBFile(reader *bufio.Reader, target []string, auth utils.SelectDB(c, lastdb) } } - if len(conf.Options.FilterKey) != 0 && !hasAtLeastOnePrefix(string(e.Key), conf.Options.FilterKey) { + if len(conf.Options.FilterKey) != 0 && !utils.HasAtLeastOnePrefix(string(e.Key), conf.Options.FilterKey) { continue } utils.RestoreRdbEntry(c, e) diff --git a/src/redis-shake/rump.go b/src/redis-shake/rump.go index 4fb6679..e22b0a4 100644 --- a/src/redis-shake/rump.go +++ b/src/redis-shake/rump.go @@ -320,7 +320,7 @@ func (dre *dbRumperExecutor) writer() { bucket := utils.StartQoS(conf.Options.Qps) preDb := 0 for ele := range dre.keyChan { - if len(conf.Options.FilterKey) != 0 && !hasAtLeastOnePrefix(ele.key, conf.Options.FilterKey) { + if len(conf.Options.FilterKey) != 0 && !utils.HasAtLeastOnePrefix(ele.key, conf.Options.FilterKey) { continue } // QoS, limit the qps diff --git a/src/redis-shake/sync.go b/src/redis-shake/sync.go index 78aac48..90f51ab 100644 --- a/src/redis-shake/sync.go +++ b/src/redis-shake/sync.go @@ -426,7 +426,7 @@ func (ds *dbSyncer) syncRDBFile(reader *bufio.Reader, target []string, auth_type } if len(conf.Options.FilterKey) != 0 { - if hasAtLeastOnePrefix(string(e.Key), conf.Options.FilterKey) { + if utils.HasAtLeastOnePrefix(string(e.Key), conf.Options.FilterKey) { utils.RestoreRdbEntry(c, e) } } else if len(conf.Options.FilterSlot) > 0 { diff --git a/src/redis-shake/utils.go b/src/redis-shake/utils.go deleted file mode 100644 index 0f21279..0000000 --- a/src/redis-shake/utils.go +++ /dev/null @@ -1,13 +0,0 @@ -package run - -import "strings" - -// hasAtLeastOnePrefix checks whether the key has begins with at least one of prefixes. -func hasAtLeastOnePrefix(key string, prefixes []string) bool { - for _, prefix := range prefixes { - if strings.HasPrefix(key, prefix) { - return true - } - } - return false -} diff --git a/src/redis-shake/utils_test.go b/src/redis-shake/utils_test.go deleted file mode 100644 index 77a3353..0000000 --- a/src/redis-shake/utils_test.go +++ /dev/null @@ -1,40 +0,0 @@ -package run - -import ( - "testing" - - "github.com/stretchr/testify/assert" -) - -func TestHasAtLeastOnePrefix(t *testing.T) { - - cases := []struct { - key string - prefixes []string - expectResult bool - }{ - { - // no prefix provided - "a", - []string{}, - false, - }, - { - // has prefix - "abc", - []string{"ab"}, - true, - }, - { - // does NOT have prefix - "abc", - []string{"edf", "wab"}, - false, - }, - } - - for _, c := range cases { - result := hasAtLeastOnePrefix(c.key, c.prefixes) - assert.Equal(t, c.expectResult, result) - } -}