From dbb83b695e8ae2cce732c1fff55bf8ab04122ab4 Mon Sep 17 00:00:00 2001 From: tjzhang-BQ <111323543+tjzhang-BQ@users.noreply.github.com> Date: Wed, 15 Jan 2025 13:43:57 -0800 Subject: [PATCH] Go: Add command ZScan (#2950) * Go: Add command ZScan Signed-off-by: TJ Zhang --- go/api/base_client.go | 87 ++++++++++++++ go/api/options/constants.go | 1 + go/api/options/zscan_options.go | 44 +++++++ go/api/sorted_set_commands.go | 4 + go/integTest/shared_commands_test.go | 165 +++++++++++++++++++++++++++ 5 files changed, 301 insertions(+) create mode 100644 go/api/options/zscan_options.go diff --git a/go/api/base_client.go b/go/api/base_client.go index de8e3f91f4..1a67892934 100644 --- a/go/api/base_client.go +++ b/go/api/base_client.go @@ -1818,3 +1818,90 @@ func (client *baseClient) ZScore(key string, member string) (Result[float64], er } return handleFloatOrNilResponse(result) } + +// Iterates incrementally over a sorted set. +// +// See [valkey.io] for details. +// +// Parameters: +// +// key - The key of the sorted set. +// cursor - The cursor that points to the next iteration of results. +// A value of `"0"` indicates the start of the search. +// For Valkey 8.0 and above, negative cursors are treated like the initial cursor("0"). +// +// Return value: +// +// The first return value is the `cursor` for the next iteration of results. `"0"` will be the `cursor` +// returned on the last iteration of the sorted set. +// The second return value is always an array of the subset of the sorted set held in `key`. +// The array is a flattened series of `string` pairs, where the value is at even indices and the score is at odd indices. +// +// Example: +// +// // assume "key" contains a set +// resCursor, resCol, err := client.ZScan("key", "0") +// fmt.Println(resCursor.Value()) +// fmt.Println(resCol.Value()) +// for resCursor != "0" { +// resCursor, resCol, err = client.ZScan("key", resCursor.Value()) +// fmt.Println("Cursor: ", resCursor.Value()) +// fmt.Println("Members: ", resCol.Value()) +// } +// +// [valkey.io]: https://valkey.io/commands/zscan/ +func (client *baseClient) ZScan(key string, cursor string) (Result[string], []Result[string], error) { + result, err := client.executeCommand(C.ZScan, []string{key, cursor}) + if err != nil { + return CreateNilStringResult(), nil, err + } + return handleScanResponse(result) +} + +// Iterates incrementally over a sorted set. +// +// See [valkey.io] for details. +// +// Parameters: +// +// key - The key of the sorted set. +// cursor - The cursor that points to the next iteration of results. +// options - The options for the command. See [options.ZScanOptions] for details. +// +// Return value: +// +// The first return value is the `cursor` for the next iteration of results. `"0"` will be the `cursor` +// returned on the last iteration of the sorted set. +// The second return value is always an array of the subset of the sorted set held in `key`. +// The array is a flattened series of `string` pairs, where the value is at even indices and the score is at odd indices. +// If `ZScanOptionsBuilder#noScores` is to `true`, the second return value will only contain the members without scores. +// +// Example: +// +// resCursor, resCol, err := client.ZScanWithOptions("key", "0", options.NewBaseScanOptionsBuilder().SetMatch("*")) +// fmt.Println(resCursor.Value()) +// fmt.Println(resCol.Value()) +// for resCursor != "0" { +// resCursor, resCol, err = client.ZScanWithOptions("key", resCursor.Value(), +// options.NewBaseScanOptionsBuilder().SetMatch("*")) +// fmt.Println("Cursor: ", resCursor.Value()) +// fmt.Println("Members: ", resCol.Value()) +// } +// +// [valkey.io]: https://valkey.io/commands/zscan/ +func (client *baseClient) ZScanWithOptions( + key string, + cursor string, + options *options.ZScanOptions, +) (Result[string], []Result[string], error) { + optionArgs, err := options.ToArgs() + if err != nil { + return CreateNilStringResult(), nil, err + } + + result, err := client.executeCommand(C.ZScan, append([]string{key, cursor}, optionArgs...)) + if err != nil { + return CreateNilStringResult(), nil, err + } + return handleScanResponse(result) +} diff --git a/go/api/options/constants.go b/go/api/options/constants.go index 83b0b3f0b8..d2d4b594db 100644 --- a/go/api/options/constants.go +++ b/go/api/options/constants.go @@ -7,4 +7,5 @@ const ( MatchKeyword string = "MATCH" // Valkey API keyword used to indicate the match filter. NoValue string = "NOVALUE" // Valkey API keyword for the no value option for hcsan command. WithScore string = "WITHSCORE" // Valkey API keyword for the with score option for zrank and zrevrank commands. + NoScores string = "NOSCORES" // Valkey API keyword for the no scores option for zscan command. ) diff --git a/go/api/options/zscan_options.go b/go/api/options/zscan_options.go new file mode 100644 index 0000000000..54fc4f3259 --- /dev/null +++ b/go/api/options/zscan_options.go @@ -0,0 +1,44 @@ +// Copyright Valkey GLIDE Project Contributors - SPDX Identifier: Apache-2.0 + +package options + +// This struct represents the optional arguments for the ZSCAN command. +type ZScanOptions struct { + BaseScanOptions + noScores bool +} + +func NewZScanOptionsBuilder() *ZScanOptions { + return &ZScanOptions{} +} + +// SetNoScores sets the noScores flag for the ZSCAN command. +// If this value is set to true, the ZSCAN command will be called with NOSCORES option. +// In the NOSCORES option, scores are not included in the response. +func (zScanOptions *ZScanOptions) SetNoScores(noScores bool) *ZScanOptions { + zScanOptions.noScores = noScores + return zScanOptions +} + +// SetMatch sets the match pattern for the ZSCAN command. +func (zScanOptions *ZScanOptions) SetMatch(match string) *ZScanOptions { + zScanOptions.BaseScanOptions.SetMatch(match) + return zScanOptions +} + +// SetCount sets the count of the ZSCAN command. +func (zScanOptions *ZScanOptions) SetCount(count int64) *ZScanOptions { + zScanOptions.BaseScanOptions.SetCount(count) + return zScanOptions +} + +func (options *ZScanOptions) ToArgs() ([]string, error) { + args := []string{} + baseArgs, err := options.BaseScanOptions.ToArgs() + args = append(args, baseArgs...) + + if options.noScores { + args = append(args, NoScores) + } + return args, err +} diff --git a/go/api/sorted_set_commands.go b/go/api/sorted_set_commands.go index 398cc7c61f..884794db9a 100644 --- a/go/api/sorted_set_commands.go +++ b/go/api/sorted_set_commands.go @@ -381,4 +381,8 @@ type SortedSetCommands interface { ZScore(key string, member string) (Result[float64], error) ZCount(key string, rangeOptions *options.ZCountRange) (int64, error) + + ZScan(key string, cursor string) (Result[string], []Result[string], error) + + ZScanWithOptions(key string, cursor string, options *options.ZScanOptions) (Result[string], []Result[string], error) } diff --git a/go/integTest/shared_commands_test.go b/go/integTest/shared_commands_test.go index 7c9473e325..ce9f7d5309 100644 --- a/go/integTest/shared_commands_test.go +++ b/go/integTest/shared_commands_test.go @@ -5043,3 +5043,168 @@ func (suite *GlideTestSuite) Test_XDel() { assert.IsType(t, &api.RequestError{}, err) }) } + +func (suite *GlideTestSuite) TestZScan() { + suite.runWithDefaultClients(func(client api.BaseClient) { + key1 := uuid.New().String() + initialCursor := "0" + defaultCount := 20 + + // Set up test data - use a large number of entries to force an iterative cursor + numberMap := make(map[string]float64) + numMembersResult := make([]api.Result[string], 50000) + charMembers := []string{"a", "b", "c", "d", "e"} + charMembersResult := []api.Result[string]{ + api.CreateStringResult("a"), + api.CreateStringResult("b"), + api.CreateStringResult("c"), + api.CreateStringResult("d"), + api.CreateStringResult("e"), + } + for i := 0; i < 50000; i++ { + numberMap["member"+strconv.Itoa(i)] = float64(i) + numMembersResult[i] = api.CreateStringResult("member" + strconv.Itoa(i)) + } + charMap := make(map[string]float64) + charMapValues := []api.Result[string]{} + for i, val := range charMembers { + charMap[val] = float64(i) + charMapValues = append(charMapValues, api.CreateStringResult(strconv.Itoa(i))) + } + + // Empty set + resCursor, resCollection, err := client.ZScan(key1, initialCursor) + assert.NoError(suite.T(), err) + assert.Equal(suite.T(), initialCursor, resCursor.Value()) + assert.Empty(suite.T(), resCollection) + + // Negative cursor + if suite.serverVersion >= "8.0.0" { + _, _, err = client.ZScan(key1, "-1") + assert.NotNil(suite.T(), err) + assert.IsType(suite.T(), &api.RequestError{}, err) + } else { + resCursor, resCollection, err = client.ZScan(key1, "-1") + assert.NoError(suite.T(), err) + assert.Equal(suite.T(), initialCursor, resCursor.Value()) + assert.Empty(suite.T(), resCollection) + } + + // Result contains the whole set + res, err := client.ZAdd(key1, charMap) + assert.NoError(suite.T(), err) + assert.Equal(suite.T(), int64(5), res) + + resCursor, resCollection, err = client.ZScan(key1, initialCursor) + assert.NoError(suite.T(), err) + assert.Equal(suite.T(), initialCursor, resCursor.Value()) + assert.Equal(suite.T(), len(charMap)*2, len(resCollection)) + + resultKeySet := make([]api.Result[string], 0, len(charMap)) + resultValueSet := make([]api.Result[string], 0, len(charMap)) + + // Iterate through array taking pairs of items + for i := 0; i < len(resCollection); i += 2 { + resultKeySet = append(resultKeySet, resCollection[i]) + resultValueSet = append(resultValueSet, resCollection[i+1]) + } + + // Verify all expected keys exist in result + assert.True(suite.T(), isSubset(charMembersResult, resultKeySet)) + + // Scores come back as integers converted to a string when the fraction is zero. + assert.True(suite.T(), isSubset(charMapValues, resultValueSet)) + + opts := options.NewZScanOptionsBuilder().SetMatch("a") + resCursor, resCollection, err = client.ZScanWithOptions(key1, initialCursor, opts) + assert.NoError(suite.T(), err) + assert.Equal(suite.T(), initialCursor, resCursor.Value()) + assert.Equal(suite.T(), resCollection, []api.Result[string]{api.CreateStringResult("a"), api.CreateStringResult("0")}) + + // Result contains a subset of the key + res, err = client.ZAdd(key1, numberMap) + assert.NoError(suite.T(), err) + assert.Equal(suite.T(), int64(50000), res) + + resCursor, resCollection, err = client.ZScan(key1, "0") + assert.NoError(suite.T(), err) + resultCollection := resCollection + resKeys := []api.Result[string]{} + + // 0 is returned for the cursor of the last iteration + for resCursor.Value() != "0" { + nextCursor, nextCol, err := client.ZScan(key1, resCursor.Value()) + assert.NoError(suite.T(), err) + assert.NotEqual(suite.T(), nextCursor, resCursor) + assert.False(suite.T(), isSubset(resultCollection, nextCol)) + resultCollection = append(resultCollection, nextCol...) + resCursor = nextCursor + } + + for i := 0; i < len(resultCollection); i += 2 { + resKeys = append(resKeys, resultCollection[i]) + } + + assert.NotEmpty(suite.T(), resultCollection) + // Verify we got all keys and values + assert.True(suite.T(), isSubset(numMembersResult, resKeys)) + + // Test match pattern + opts = options.NewZScanOptionsBuilder().SetMatch("*") + resCursor, resCollection, err = client.ZScanWithOptions(key1, initialCursor, opts) + assert.NoError(suite.T(), err) + assert.NotEqual(suite.T(), initialCursor, resCursor.Value()) + assert.GreaterOrEqual(suite.T(), len(resCollection), defaultCount) + + // test count + opts = options.NewZScanOptionsBuilder().SetCount(20) + resCursor, resCollection, err = client.ZScanWithOptions(key1, initialCursor, opts) + assert.NoError(suite.T(), err) + assert.NotEqual(suite.T(), initialCursor, resCursor.Value()) + assert.GreaterOrEqual(suite.T(), len(resCollection), 20) + + // test count with match, returns a non-empty array + opts = options.NewZScanOptionsBuilder().SetMatch("1*").SetCount(20) + resCursor, resCollection, err = client.ZScanWithOptions(key1, initialCursor, opts) + assert.NoError(suite.T(), err) + assert.NotEqual(suite.T(), initialCursor, resCursor.Value()) + assert.GreaterOrEqual(suite.T(), len(resCollection), 0) + + // Test NoScores option for Redis 8.0.0+ + if suite.serverVersion >= "8.0.0" { + opts = options.NewZScanOptionsBuilder().SetNoScores(true) + resCursor, resCollection, err = client.ZScanWithOptions(key1, initialCursor, opts) + assert.NoError(suite.T(), err) + cursor, err := strconv.ParseInt(resCursor.Value(), 10, 64) + assert.NoError(suite.T(), err) + assert.GreaterOrEqual(suite.T(), cursor, int64(0)) + + // Verify all fields start with "member" + for _, field := range resCollection { + assert.True(suite.T(), strings.HasPrefix(field.Value(), "member")) + } + } + + // Test exceptions + // Non-set key + stringKey := uuid.New().String() + setRes, err := client.Set(stringKey, "test") + assert.NoError(suite.T(), err) + assert.Equal(suite.T(), "OK", setRes) + + _, _, err = client.ZScan(stringKey, initialCursor) + assert.NotNil(suite.T(), err) + assert.IsType(suite.T(), &api.RequestError{}, err) + + opts = options.NewZScanOptionsBuilder().SetMatch("test").SetCount(1) + _, _, err = client.ZScanWithOptions(stringKey, initialCursor, opts) + assert.NotNil(suite.T(), err) + assert.IsType(suite.T(), &api.RequestError{}, err) + + // Negative count + opts = options.NewZScanOptionsBuilder().SetCount(-1) + _, _, err = client.ZScanWithOptions(key1, "-1", opts) + assert.NotNil(suite.T(), err) + assert.IsType(suite.T(), &api.RequestError{}, err) + }) +}