From fe9f9eda0a7c54819c79d6169e8c9308b9d59110 Mon Sep 17 00:00:00 2001 From: tjzhang-BQ <111323543+tjzhang-BQ@users.noreply.github.com> Date: Mon, 23 Dec 2024 16:48:33 -0800 Subject: [PATCH] Go: Add ZPopMin and ZPopMax (#2850) * Go: Add ZPopMin and ZPopMax Signed-off-by: TJ Zhang --- CHANGELOG.md | 1 + go/api/base_client.go | 32 +++++++++++ go/api/response_handlers.go | 24 ++++++++ go/api/sorted_set_commands.go | 82 ++++++++++++++++++++++++++++ go/integTest/shared_commands_test.go | 69 +++++++++++++++++++++++ 5 files changed, 208 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 43148b610e..f20efd03bc 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -11,6 +11,7 @@ * Go: Add `ZIncrBy` command ([#2830](https://github.com/valkey-io/valkey-glide/pull/2830)) * Go: Add `SScan` and `SMove` ([#2789](https://github.com/valkey-io/valkey-glide/issues/2789)) * Go: Add `ZADD` ([#2813](https://github.com/valkey-io/valkey-glide/issues/2813)) +* Go: Add `ZPopMin` and `ZPopMax` ([#2850](https://github.com/valkey-io/valkey-glide/pull/2850)) #### Breaking Changes diff --git a/go/api/base_client.go b/go/api/base_client.go index 6086a8f880..7e1f5722a1 100644 --- a/go/api/base_client.go +++ b/go/api/base_client.go @@ -1316,3 +1316,35 @@ func (client *baseClient) ZIncrBy(key string, increment float64, member string) return handleDoubleResponse(result) } + +func (client *baseClient) ZPopMin(key string) (map[Result[string]]Result[float64], error) { + result, err := client.executeCommand(C.ZPopMin, []string{key}) + if err != nil { + return nil, err + } + return handleStringDoubleMapResponse(result) +} + +func (client *baseClient) ZPopMinWithCount(key string, count int64) (map[Result[string]]Result[float64], error) { + result, err := client.executeCommand(C.ZPopMin, []string{key, utils.IntToString(count)}) + if err != nil { + return nil, err + } + return handleStringDoubleMapResponse(result) +} + +func (client *baseClient) ZPopMax(key string) (map[Result[string]]Result[float64], error) { + result, err := client.executeCommand(C.ZPopMax, []string{key}) + if err != nil { + return nil, err + } + return handleStringDoubleMapResponse(result) +} + +func (client *baseClient) ZPopMaxWithCount(key string, count int64) (map[Result[string]]Result[float64], error) { + result, err := client.executeCommand(C.ZPopMax, []string{key, utils.IntToString(count)}) + if err != nil { + return nil, err + } + return handleStringDoubleMapResponse(result) +} diff --git a/go/api/response_handlers.go b/go/api/response_handlers.go index 618a95bfc2..dd4c2d1f24 100644 --- a/go/api/response_handlers.go +++ b/go/api/response_handlers.go @@ -299,6 +299,30 @@ func handleBooleanArrayResponse(response *C.struct_CommandResponse) ([]Result[bo return slice, nil } +func handleStringDoubleMapResponse(response *C.struct_CommandResponse) (map[Result[string]]Result[float64], error) { + defer C.free_command_response(response) + + typeErr := checkResponseType(response, C.Map, false) + if typeErr != nil { + return nil, typeErr + } + + m := make(map[Result[string]]Result[float64], response.array_value_len) + for _, v := range unsafe.Slice(response.array_value, response.array_value_len) { + key, err := convertCharArrayToString(v.map_key, true) + if err != nil { + return nil, err + } + typeErr := checkResponseType(v.map_value, C.Float, false) + if typeErr != nil { + return nil, typeErr + } + value := CreateFloat64Result(float64(v.map_value.float_value)) + m[key] = value + } + return m, nil +} + func handleStringToStringMapResponse(response *C.struct_CommandResponse) (map[Result[string]]Result[string], error) { defer C.free_command_response(response) diff --git a/go/api/sorted_set_commands.go b/go/api/sorted_set_commands.go index 02a9697e6d..0296dae8d5 100644 --- a/go/api/sorted_set_commands.go +++ b/go/api/sorted_set_commands.go @@ -110,4 +110,86 @@ type SortedSetCommands interface { // // [valkey.io]: https://valkey.io/commands/zincrby/ ZIncrBy(key string, increment float64, member string) (Result[float64], error) + + // Removes and returns the member with the lowest score from the sorted set + // stored at the specified `key`. + // + // see [valkey.io] for details. + // + // Parameters: + // key - The key of the sorted set. + // + // Return value: + // A map containing the removed member and its corresponding score. + // If `key` doesn't exist, it will be treated as an empty sorted set and the + // command returns an empty map. + // + // Example: + // res, err := client.zpopmin("mySortedSet") + // fmt.Println(res.Value()) // Output: map["member1":5.0] + // + // [valkey.io]: https://valkey.io/commands/zpopmin/ + ZPopMin(key string) (map[Result[string]]Result[float64], error) + + // Removes and returns up to `count` members with the lowest scores from the sorted set + // stored at the specified `key`. + // + // see [valkey.io] for details. + // + // Parameters: + // key - The key of the sorted set. + // count - The number of members to remove. + // + // Return value: + // A map containing the removed members and their corresponding scores. + // If `key` doesn't exist, it will be treated as an empty sorted set and the + // command returns an empty map. + // + // Example: + // res, err := client.ZPopMinWithCount("mySortedSet", 2) + // fmt.Println(res.Value()) // Output: map["member1":5.0, "member2":6.0] + // + // [valkey.io]: https://valkey.io/commands/zpopmin/ + ZPopMinWithCount(key string, count int64) (map[Result[string]]Result[float64], error) + + // Removes and returns the member with the highest score from the sorted set stored at the + // specified `key`. + // + // see [valkey.io] for details. + // + // Parameters: + // key - The key of the sorted set. + // + // Return value: + // A map containing the removed member and its corresponding score. + // If `key` doesn't exist, it will be treated as an empty sorted set and the + // command returns an empty map. + // + // Example: + // res, err := client.zpopmax("mySortedSet") + // fmt.Println(res.Value()) // Output: map["member2":8.0] + // + // [valkey.io]: https://valkey.io/commands/zpopmin/ + ZPopMax(key string) (map[Result[string]]Result[float64], error) + + // Removes and returns up to `count` members with the highest scores from the sorted set + // stored at the specified `key`. + // + // see [valkey.io] for details. + // + // Parameters: + // key - The key of the sorted set. + // count - The number of members to remove. + // + // Return value: + // A map containing the removed members and their corresponding scores. + // If `key` doesn't exist, it will be treated as an empty sorted set and the + // command returns an empty map. + // + // Example: + // res, err := client.ZPopMaxWithCount("mySortedSet", 2) + // fmt.Println(res.Value()) // Output: map["member1":5.0, "member2":6.0] + // + // [valkey.io]: https://valkey.io/commands/zpopmin/ + ZPopMaxWithCount(key string, count int64) (map[Result[string]]Result[float64], error) } diff --git a/go/integTest/shared_commands_test.go b/go/integTest/shared_commands_test.go index d0902439fb..c812f93bcc 100644 --- a/go/integTest/shared_commands_test.go +++ b/go/integTest/shared_commands_test.go @@ -3970,3 +3970,72 @@ func (suite *GlideTestSuite) TestZincrBy() { assert.IsType(suite.T(), &api.RequestError{}, err) }) } + +func (suite *GlideTestSuite) TestZPopMin() { + suite.runWithDefaultClients(func(client api.BaseClient) { + key1 := uuid.New().String() + key2 := uuid.New().String() + memberScoreMap := map[string]float64{ + "one": 1.0, + "two": 2.0, + "three": 3.0, + } + + res, err := client.ZAdd(key1, memberScoreMap) + assert.Nil(suite.T(), err) + assert.Equal(suite.T(), int64(3), res.Value()) + + res2, err := client.ZPopMin(key1) + assert.Nil(suite.T(), err) + assert.Len(suite.T(), res2, 1) + assert.Equal(suite.T(), float64(1.0), res2[api.CreateStringResult("one")].Value()) + + res3, err := client.ZPopMinWithCount(key1, 2) + assert.Nil(suite.T(), err) + assert.Len(suite.T(), res3, 2) + assert.Equal(suite.T(), float64(2.0), res3[api.CreateStringResult("two")].Value()) + assert.Equal(suite.T(), float64(3.0), res3[api.CreateStringResult("three")].Value()) + + // non sorted set key + _, err = client.Set(key2, "test") + assert.Nil(suite.T(), err) + + _, err = client.ZPopMin(key2) + assert.NotNil(suite.T(), err) + assert.IsType(suite.T(), &api.RequestError{}, err) + }) +} + +func (suite *GlideTestSuite) TestZPopMax() { + suite.runWithDefaultClients(func(client api.BaseClient) { + key1 := uuid.New().String() + key2 := uuid.New().String() + memberScoreMap := map[string]float64{ + "one": 1.0, + "two": 2.0, + "three": 3.0, + } + res, err := client.ZAdd(key1, memberScoreMap) + assert.Nil(suite.T(), err) + assert.Equal(suite.T(), int64(3), res.Value()) + + res2, err := client.ZPopMax(key1) + assert.Nil(suite.T(), err) + assert.Len(suite.T(), res2, 1) + assert.Equal(suite.T(), float64(3.0), res2[api.CreateStringResult("three")].Value()) + + res3, err := client.ZPopMaxWithCount(key1, 2) + assert.Nil(suite.T(), err) + assert.Len(suite.T(), res3, 2) + assert.Equal(suite.T(), float64(2.0), res3[api.CreateStringResult("two")].Value()) + assert.Equal(suite.T(), float64(1.0), res3[api.CreateStringResult("one")].Value()) + + // non sorted set key + _, err = client.Set(key2, "test") + assert.Nil(suite.T(), err) + + _, err = client.ZPopMax(key2) + assert.NotNil(suite.T(), err) + assert.IsType(suite.T(), &api.RequestError{}, err) + }) +}