Merge pull request #62 from KyberNetwork/ft/redis-scan

ft: utils to hscan only keys or values

Author: NgoKimPhu

pkg/client/redis.go

@@ -10,7 +10,7 @@ import (
 	"github.com/redis/go-redis/extra/redisotel/v9"
 	"github.com/redis/go-redis/v9"
 
-	"github.com/KyberNetwork/service-framework/pkg/client/redis/reconnectable"
+	reconredis "github.com/KyberNetwork/service-framework/pkg/client/redis/reconnectable"
 )
 
 const RedisCloseDelay = time.Minute
@@ -41,24 +41,24 @@ func (*RedisCfg) OnUpdate(old, new *RedisCfg) {
 func NewRedisClient(ctx context.Context, opts *redis.UniversalOptions) redis.UniversalClient {
 	if opts.MasterName == "" {
 		return reconredis.New(func() redis.UniversalClient {
-			return instrumentRedisOtel(ctx, redis.NewUniversalClient(opts))
+			return InstrumentRedisOtel(ctx, redis.NewUniversalClient(opts))
 		})
 	}
 	failoverOpts := opts.Failover()
 	failoverOpts.RouteByLatency = opts.RouteByLatency
 	failoverOpts.RouteRandomly = opts.RouteRandomly
-	return instrumentRedisOtel(ctx, redis.NewFailoverClusterClient(failoverOpts))
+	return InstrumentRedisOtel(ctx, redis.NewFailoverClusterClient(failoverOpts))
 }
 
-func instrumentRedisOtel(ctx context.Context, client redis.UniversalClient) redis.UniversalClient {
+func InstrumentRedisOtel(ctx context.Context, client redis.UniversalClient) redis.UniversalClient {
 	if metric.Provider() != nil {
 		if err := redisotel.InstrumentMetrics(client); err != nil {
-			klog.Errorf(ctx, "instrumentRedisOtel|redisotel.InstrumentMetrics failed|err=%v", err)
+			klog.Errorf(ctx, "InstrumentRedisOtel|redisotel.InstrumentMetrics failed|err=%v", err)
 		}
 	}
 	if tracer.Provider() != nil {
 		if err := redisotel.InstrumentTracing(client); err != nil {
-			klog.Errorf(ctx, "instrumentRedisOtel|redisotel.InstrumentTracing failed|err=%v", err)
+			klog.Errorf(ctx, "InstrumentRedisOtel|redisotel.InstrumentTracing failed|err=%v", err)
 		}
 	}
 	return client

pkg/client/redis/hscan.go

@@ -0,0 +1,57 @@
+package kredis
+
+import (
+	"context"
+
+	"github.com/redis/go-redis/v9"
+)
+
+const LuaHscanOnlyKeys = `
+	local results = {}
+	local scan_result = redis.call('HSCAN', KEYS[1], unpack(ARGV))
+	for i = 1, #scan_result[2], 2 do
+		table.insert(results, scan_result[2][i])
+	end
+	return {scan_result[1], results}
+`
+
+// HscanOnlyKeys scans only keys (no values) from a redis hash
+func HscanOnlyKeys(ctx context.Context, client redis.UniversalClient, key string, cursor uint64, match string,
+	count int64) *redis.ScanCmd {
+	args := []any{"eval", LuaHscanOnlyKeys, 1, key, cursor}
+	if match != "" {
+		args = append(args, "match", match)
+	}
+	if count > 0 {
+		args = append(args, "count", count)
+	}
+	c := client.Process
+	cmd := redis.NewScanCmd(ctx, c, args...)
+	_ = c(ctx, cmd)
+	return cmd
+}
+
+const LuaHscanOnlyValues = `
+	local results = {}
+	local scan_result = redis.call('HSCAN', KEYS[1], unpack(ARGV))
+	for i = 2, #scan_result[2], 2 do
+		table.insert(results, scan_result[2][i])
+	end
+	return {scan_result[1], results}
+`
+
+// HscanOnlyValues scans only values (no keys) from a redis hash
+func HscanOnlyValues(ctx context.Context, client redis.UniversalClient, key string, cursor uint64, match string,
+	count int64) *redis.ScanCmd {
+	args := []any{"eval", LuaHscanOnlyValues, 1, key, cursor}
+	if match != "" {
+		args = append(args, "match", match)
+	}
+	if count > 0 {
+		args = append(args, "count", count)
+	}
+	c := client.Process
+	cmd := redis.NewScanCmd(ctx, c, args...)
+	_ = c(ctx, cmd)
+	return cmd
+}

pkg/client/redis/hscan_test.go

@@ -0,0 +1,70 @@
+package kredis
+
+import (
+	"context"
+	"testing"
+
+	"github.com/alicebob/miniredis/v2"
+	"github.com/redis/go-redis/v9"
+	"github.com/stretchr/testify/suite"
+)
+
+func TestTestSuite(t *testing.T) {
+	suite.Run(t, new(TestSuite))
+}
+
+type TestSuite struct {
+	suite.Suite
+	mockRedis *miniredis.Miniredis
+	client    redis.UniversalClient
+}
+
+func (ts *TestSuite) SetupSuite() {
+	ts.mockRedis = miniredis.RunT(ts.T())
+	ts.client = redis.NewClient(&redis.Options{
+		Addr: ts.mockRedis.Addr(),
+	})
+
+	// Populate mock Redis with some data
+	ts.mockRedis.HSet("myhash", "key1", "value1", "key2", "value2", "key3", "value3")
+}
+
+func (ts *TestSuite) TestHscanOnlyKeys() {
+	ctx := context.Background()
+
+	// Test HscanOnlyKeys
+	scanCmd := HscanOnlyKeys(ctx, ts.client, "myhash", 0, "", 0)
+	keys, cursor, err := scanCmd.Result()
+
+	ts.Nil(err)
+	ts.Equal(uint64(0), cursor)
+	ts.ElementsMatch([]string{"key1", "key2", "key3"}, keys)
+
+	// Test with match pattern
+	scanCmd = HscanOnlyKeys(ctx, ts.client, "myhash", 0, "key[1-2]", 0)
+	keys, cursor, err = scanCmd.Result()
+
+	ts.Nil(err)
+	ts.Equal(uint64(0), cursor)
+	ts.ElementsMatch([]string{"key1", "key2"}, keys)
+}
+
+func (ts *TestSuite) TestHscanOnlyValues() {
+	ctx := context.Background()
+
+	// Test HscanOnlyValues
+	scanCmd := HscanOnlyValues(ctx, ts.client, "myhash", 0, "", 0)
+	values, cursor, err := scanCmd.Result()
+
+	ts.Nil(err)
+	ts.Equal(uint64(0), cursor)
+	ts.ElementsMatch([]string{"value1", "value2", "value3"}, values)
+
+	// Test with match pattern
+	scanCmd = HscanOnlyValues(ctx, ts.client, "myhash", 0, "key[1-2]", 0)
+	values, cursor, err = scanCmd.Result()
+
+	ts.Nil(err)
+	ts.Equal(uint64(0), cursor)
+	ts.ElementsMatch([]string{"value1", "value2"}, values)
+}

pkg/client/redis/paginate.go

@@ -0,0 +1,25 @@
+package kredis
+
+import (
+	"errors"
+
+	"github.com/redis/go-redis/v9"
+)
+
+func PaginateScanCmd(scanFn func(cursor uint64) *redis.ScanCmd, callbackFn func(keys []string)) (err error) {
+	var keys []string
+	for cursor := uint64(0); ; {
+		if keys, cursor, err = scanFn(cursor).Result(); err != nil {
+			if errors.Is(err, redis.Nil) {
+				err = nil
+			}
+			return err
+		}
+
+		callbackFn(keys)
+
+		if cursor == 0 {
+			return nil
+		}
+	}
+}