From a3b3da6ab150b654c5c0bfd0ed14b39cd74c5e0f Mon Sep 17 00:00:00 2001 From: tjzhang-BQ <111323543+tjzhang-BQ@users.noreply.github.com> Date: Fri, 10 Jan 2025 10:56:47 -0800 Subject: [PATCH] Go: Add command ZRank and ZRevRank (#2932) * Go: Add command ZRank and ZRevRank Signed-off-by: TJ Zhang --- go/api/base_client.go | 32 ++++++++ go/api/options/constants.go | 7 +- go/api/response_handlers.go | 26 +++++++ go/api/sorted_set_commands.go | 110 +++++++++++++++++++++++++++ go/integTest/shared_commands_test.go | 70 +++++++++++++++++ 5 files changed, 242 insertions(+), 3 deletions(-) diff --git a/go/api/base_client.go b/go/api/base_client.go index be3b98c198..a9d8237911 100644 --- a/go/api/base_client.go +++ b/go/api/base_client.go @@ -1503,3 +1503,35 @@ func (client *baseClient) XLen(key string) (Result[int64], error) { } return handleLongResponse(result) } + +func (client *baseClient) ZRank(key string, member string) (Result[int64], error) { + result, err := client.executeCommand(C.ZRank, []string{key, member}) + if err != nil { + return CreateNilInt64Result(), err + } + return handleLongOrNullResponse(result) +} + +func (client *baseClient) ZRankWithScore(key string, member string) (Result[int64], Result[float64], error) { + result, err := client.executeCommand(C.ZRank, []string{key, member, options.WithScore}) + if err != nil { + return CreateNilInt64Result(), CreateNilFloat64Result(), err + } + return handleLongAndDoubleOrNullResponse(result) +} + +func (client *baseClient) ZRevRank(key string, member string) (Result[int64], error) { + result, err := client.executeCommand(C.ZRevRank, []string{key, member}) + if err != nil { + return CreateNilInt64Result(), err + } + return handleLongOrNullResponse(result) +} + +func (client *baseClient) ZRevRankWithScore(key string, member string) (Result[int64], Result[float64], error) { + result, err := client.executeCommand(C.ZRevRank, []string{key, member, options.WithScore}) + if err != nil { + return CreateNilInt64Result(), CreateNilFloat64Result(), err + } + return handleLongAndDoubleOrNullResponse(result) +} diff --git a/go/api/options/constants.go b/go/api/options/constants.go index f38b0f4541..83b0b3f0b8 100644 --- a/go/api/options/constants.go +++ b/go/api/options/constants.go @@ -3,7 +3,8 @@ package options const ( - CountKeyword string = "COUNT" // Valkey API keyword used to extract specific number of matching indices from a list. - 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. + CountKeyword string = "COUNT" // Valkey API keyword used to extract specific number of matching indices from a list. + 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. ) diff --git a/go/api/response_handlers.go b/go/api/response_handlers.go index 9f788f507d..fe2ecde613 100644 --- a/go/api/response_handlers.go +++ b/go/api/response_handlers.go @@ -267,6 +267,32 @@ func handleDoubleResponse(response *C.struct_CommandResponse) (Result[float64], return CreateFloat64Result(float64(response.float_value)), nil } +func handleLongAndDoubleOrNullResponse(response *C.struct_CommandResponse) (Result[int64], Result[float64], error) { + defer C.free_command_response(response) + + typeErr := checkResponseType(response, C.Array, true) + if typeErr != nil { + return CreateNilInt64Result(), CreateNilFloat64Result(), typeErr + } + + if response.response_type == C.Null { + return CreateNilInt64Result(), CreateNilFloat64Result(), nil + } + + rank := CreateNilInt64Result() + score := CreateNilFloat64Result() + for _, v := range unsafe.Slice(response.array_value, response.array_value_len) { + if v.response_type == C.Int { + rank = CreateInt64Result(int64(v.int_value)) + } + if v.response_type == C.Float { + score = CreateFloat64Result(float64(v.float_value)) + } + } + + return rank, score, nil +} + func handleBooleanResponse(response *C.struct_CommandResponse) (Result[bool], error) { defer C.free_command_response(response) diff --git a/go/api/sorted_set_commands.go b/go/api/sorted_set_commands.go index 4b63b70091..510a28a3fc 100644 --- a/go/api/sorted_set_commands.go +++ b/go/api/sorted_set_commands.go @@ -262,4 +262,114 @@ type SortedSetCommands interface { // [valkey.io]: https://valkey.io/commands/bzpopmin/ // [blocking commands]: https://github.com/valkey-io/valkey-glide/wiki/General-Concepts#blocking-commands BZPopMin(keys []string, timeoutSecs float64) (Result[KeyWithMemberAndScore], error) + + // Returns the rank of `member` in the sorted set stored at `key`, with + // scores ordered from low to high, starting from `0`. + // To get the rank of `member` with its score, see [ZRankWithScore]. + // + // See [valkey.io] for details. + // + // Parameters: + // key - The key of the sorted set. + // member - The member to get the rank of. + // + // Return value: + // The rank of `member` in the sorted set. + // If `key` doesn't exist, or if `member` is not present in the set, + // `nil` will be returned. + // + // Example: + // res, err := client.ZRank("mySortedSet", "member1") + // fmt.Println(res.Value()) // Output: 3 + // + // res2, err := client.ZRank("mySortedSet", "non-existing-member") + // if res2.IsNil() { + // fmt.Println("Member not found") + // } + // + // [valkey.io]: https://valkey.io/commands/zrank/ + ZRank(key string, member string) (Result[int64], error) + + // Returns the rank of `member` in the sorted set stored at `key` with its + // score, where scores are ordered from the lowest to highest, starting from `0`. + // + // See [valkey.io] for details. + // + // Parameters: + // key - The key of the sorted set. + // member - The member to get the rank of. + // + // Return value: + // A tuple containing the rank of `member` and its score. + // If `key` doesn't exist, or if `member` is not present in the set, + // `nil` will be returned. + // + // Example: + // resRank, resScore, err := client.ZRankWithScore("mySortedSet", "member1") + // fmt.Println(resRank.Value()) // Output: 3 + // fmt.Println(resScore.Value()) // Output: 5.0 + // + // res2Rank, res2Score, err := client.ZRankWithScore("mySortedSet", "non-existing-member") + // if res2Rank.IsNil() { + // fmt.Println("Member not found") + // } + // + // [valkey.io]: https://valkey.io/commands/zrank/ + ZRankWithScore(key string, member string) (Result[int64], Result[float64], error) + + // Returns the rank of `member` in the sorted set stored at `key`, where + // scores are ordered from the highest to lowest, starting from `0`. + // To get the rank of `member` with its score, see [ZRevRankWithScore]. + // + // See [valkey.io] for details. + // + // Parameters: + // key - The key of the sorted set. + // member - The member to get the rank of. + // + // Return value: + // The rank of `member` in the sorted set, where ranks are ordered from high to + // low based on scores. + // If `key` doesn't exist, or if `member` is not present in the set, + // `nil` will be returned. + // + // Example: + // res, err := client.ZRevRank("mySortedSet", "member2") + // fmt.Println(res.Value()) // Output: 1 + // + // res2, err := client.ZRevRank("mySortedSet", "non-existing-member") + // if res2.IsNil() { + // fmt.Println("Member not found") + // } + // + // [valkey.io]: https://valkey.io/commands/zrevrank/ + ZRevRank(key string, member string) (Result[int64], error) + + // Returns the rank of `member` in the sorted set stored at `key`, where + // scores are ordered from the highest to lowest, starting from `0`. + // To get the rank of `member` with its score, see [ZRevRankWithScore]. + // + // See [valkey.io] for details. + // + // Parameters: + // key - The key of the sorted set. + // member - The member to get the rank of. + // + // Return value: + // A tuple containing the rank of `member` and its score. + // If `key` doesn't exist, or if `member` is not present in the set, + // `nil` will be returned.s + // + // Example: + // resRank, resScore, err := client.ZRevRankWithScore("mySortedSet", "member2") + // fmt.Println(resRank.Value()) // Output: 1 + // fmt.Println(resScore.Value()) // Output: 6.0 + // + // res2Rank, res2Score, err := client.ZRevRankWithScore("mySortedSet", "non-existing-member") + // if res2Rank.IsNil() { + // fmt.Println("Member not found") + // } + // + // [valkey.io]: https://valkey.io/commands/zrevrank/ + ZRevRankWithScore(key string, member string) (Result[int64], Result[float64], error) } diff --git a/go/integTest/shared_commands_test.go b/go/integTest/shared_commands_test.go index 4084e4991d..da0002d5c6 100644 --- a/go/integTest/shared_commands_test.go +++ b/go/integTest/shared_commands_test.go @@ -4578,3 +4578,73 @@ func (suite *GlideTestSuite) Test_XAdd_XLen_XTrim() { assert.IsType(t, &api.RequestError{}, err) }) } + +func (suite *GlideTestSuite) TestZRank() { + suite.runWithDefaultClients(func(client api.BaseClient) { + key := uuid.New().String() + stringKey := uuid.New().String() + client.ZAdd(key, map[string]float64{"one": 1.5, "two": 2.0, "three": 3.0}) + res, err := client.ZRank(key, "two") + assert.Nil(suite.T(), err) + assert.Equal(suite.T(), int64(1), res.Value()) + + if suite.serverVersion >= "7.2.0" { + res2Rank, res2Score, err := client.ZRankWithScore(key, "one") + assert.Nil(suite.T(), err) + assert.Equal(suite.T(), int64(0), res2Rank.Value()) + assert.Equal(suite.T(), float64(1.5), res2Score.Value()) + res4Rank, res4Score, err := client.ZRankWithScore(key, "non-existing-member") + assert.Nil(suite.T(), err) + assert.True(suite.T(), res4Rank.IsNil()) + assert.True(suite.T(), res4Score.IsNil()) + } + + res3, err := client.ZRank(key, "non-existing-member") + assert.Nil(suite.T(), err) + assert.True(suite.T(), res3.IsNil()) + + // key exists, but it is not a set + setRes, err := client.Set(stringKey, "value") + assert.Nil(suite.T(), err) + assert.Equal(suite.T(), "OK", setRes.Value()) + + _, err = client.ZRank(stringKey, "value") + assert.NotNil(suite.T(), err) + assert.IsType(suite.T(), &api.RequestError{}, err) + }) +} + +func (suite *GlideTestSuite) TestZRevRank() { + suite.runWithDefaultClients(func(client api.BaseClient) { + key := uuid.New().String() + stringKey := uuid.New().String() + client.ZAdd(key, map[string]float64{"one": 1.5, "two": 2.0, "three": 3.0}) + res, err := client.ZRevRank(key, "two") + assert.Nil(suite.T(), err) + assert.Equal(suite.T(), int64(1), res.Value()) + + if suite.serverVersion >= "7.2.0" { + res2Rank, res2Score, err := client.ZRevRankWithScore(key, "one") + assert.Nil(suite.T(), err) + assert.Equal(suite.T(), int64(2), res2Rank.Value()) + assert.Equal(suite.T(), float64(1.5), res2Score.Value()) + res4Rank, res4Score, err := client.ZRevRankWithScore(key, "non-existing-member") + assert.Nil(suite.T(), err) + assert.True(suite.T(), res4Rank.IsNil()) + assert.True(suite.T(), res4Score.IsNil()) + } + + res3, err := client.ZRevRank(key, "non-existing-member") + assert.Nil(suite.T(), err) + assert.True(suite.T(), res3.IsNil()) + + // key exists, but it is not a set + setRes, err := client.Set(stringKey, "value") + assert.Nil(suite.T(), err) + assert.Equal(suite.T(), "OK", setRes.Value()) + + _, err = client.ZRevRank(stringKey, "value") + assert.NotNil(suite.T(), err) + assert.IsType(suite.T(), &api.RequestError{}, err) + }) +}