From 81f8dd5bc18653ee5762bb8db1ae4397b9e3f0eb Mon Sep 17 00:00:00 2001 From: TJ Zhang Date: Thu, 16 Jan 2025 19:26:22 -0800 Subject: [PATCH] Go: Add commands ZRemRangeByRank/ZRemRangeByScore/ZRemRangeByLex Signed-off-by: TJ Zhang --- go/api/base_client.go | 95 ++++++++++++++++ go/api/options/zrange_options.go | 12 ++ go/api/sorted_set_commands.go | 6 + go/integTest/shared_commands_test.go | 159 +++++++++++++++++++++++++++ 4 files changed, 272 insertions(+) diff --git a/go/api/base_client.go b/go/api/base_client.go index 5b27cf2d11..a3c93f97ee 100644 --- a/go/api/base_client.go +++ b/go/api/base_client.go @@ -2183,3 +2183,98 @@ func (client *baseClient) Echo(message string) (Result[string], error) { } return handleStringOrNilResponse(result) } + +// Removes all elements in the sorted set stored at `key` with a lexicographical order +// between `rangeQuery.Start` and `rangeQuery.End`. +// +// See [valkey.io] for details. +// +// Parameters: +// +// key - The key of the sorted set. +// rangeQuery - The range query object representing the minimum and maximum bound of the lexicographical range. +// can be an implementation of [options.LexBoundary]. +// +// Return value: +// +// The number of members removed from the sorted set. +// If `key` does not exist, it is treated as an empty sorted set, and the command returns `0`. +// If `rangeQuery.Start` is greater than `rangeQuery.End`, `0` is returned. +// +// Example: +// +// zRemRangeByLexResult, err := client.ZRemRangeByLex("key1", options.NewRangeByLexQuery("a", "b")) +// fmt.Println(zRemRangeByLexResult) // Output: 1 +// +// [valkey.io]: https://valkey.io/commands/zremrangebylex/ +func (client *baseClient) ZRemRangeByLex(key string, rangeQuery options.RangeByLex) (int64, error) { + result, err := client.executeCommand( + C.ZRemRangeByLex, append([]string{key}, rangeQuery.ToArgsRemRange()...)) + if err != nil { + return defaultIntResponse, err + } + return handleIntResponse(result) +} + +// Removes all elements in the sorted set stored at `key` with a rank between `start` and `stop`. +// +// See [valkey.io] for details. +// +// Parameters: +// +// key - The key of the sorted set. +// start - The start rank. +// stop - The stop rank. +// +// Return value: +// +// The number of members removed from the sorted set. +// If `key` does not exist, it is treated as an empty sorted set, and the command returns `0`. +// If `start` is greater than `stop`, `0` is returned. +// +// Example: +// +// zRemRangeByRankResult, err := client.ZRemRangeByRank("key1", 0, 1) +// fmt.Println(zRemRangeByRankResult) // Output: 1 +// +// [valkey.io]: https://valkey.io/commands/zremrangebyrank/ +func (client *baseClient) ZRemRangeByRank(key string, start int64, stop int64) (int64, error) { + result, err := client.executeCommand(C.ZRemRangeByRank, []string{key, utils.IntToString(start), utils.IntToString(stop)}) + if err != nil { + return 0, err + } + return handleIntResponse(result) +} + +// Removes all elements in the sorted set stored at `key` with a score between `rangeQuery.Start` and `rangeQuery.End`. +// +// See [valkey.io] for details. +// +// Parameters: +// +// key - The key of the sorted set. +// rangeQuery - The range query object representing the minimum and maximum bound of the score range. +// can be an implementation of [options.RangeByScore]. +// +// Return value: +// +// The number of members removed from the sorted set. +// If `key` does not exist, it is treated as an empty sorted set, and the command returns `0`. +// If `rangeQuery.Start` is greater than `rangeQuery.End`, `0` is returned. +// +// Example: +// +// zRemRangeByScoreResult, err := client.ZRemRangeByScore("key1", options.NewRangeByScoreBuilder( +// options.NewInfiniteScoreBoundary(options.NegativeInfinity), +// options.NewInfiniteScoreBoundary(options.PositiveInfinity), +// )) +// fmt.Println(zRemRangeByScoreResult) // Output: 1 +// +// [valkey.io]: https://valkey.io/commands/zremrangebyscore/ +func (client *baseClient) ZRemRangeByScore(key string, rangeQuery options.RangeByScore) (int64, error) { + result, err := client.executeCommand(C.ZRemRangeByScore, append([]string{key}, rangeQuery.ToArgsRemRange()...)) + if err != nil { + return 0, err + } + return handleIntResponse(result) +} diff --git a/go/api/options/zrange_options.go b/go/api/options/zrange_options.go index 002dc38e24..d89e9124b8 100644 --- a/go/api/options/zrange_options.go +++ b/go/api/options/zrange_options.go @@ -14,6 +14,10 @@ type ZRangeQuery interface { ToArgs() []string } +type ZRemRangeQuery interface { + ToArgsRemRange() []string +} + // Queries a range of elements from a sorted set by theirs index. type RangeByIndex struct { start, end int64 @@ -152,6 +156,10 @@ func (rbs *RangeByScore) ToArgs() []string { return args } +func (rbs *RangeByScore) ToArgsRemRange() []string { + return []string{string(rbs.start), string(rbs.end)} +} + // Queries a range of elements from a sorted set by theirs lexicographical order. // // Parameters: @@ -186,6 +194,10 @@ func (rbl *RangeByLex) ToArgs() []string { return args } +func (rbl *RangeByLex) ToArgsRemRange() []string { + return []string{string(rbl.start), string(rbl.end)} +} + // Query for `ZRangeWithScores` in [SortedSetCommands] // - For range queries by index (rank), use `RangeByIndex`. // - For range queries by score, use `RangeByScore`. diff --git a/go/api/sorted_set_commands.go b/go/api/sorted_set_commands.go index 884794db9a..237a8d3489 100644 --- a/go/api/sorted_set_commands.go +++ b/go/api/sorted_set_commands.go @@ -385,4 +385,10 @@ type SortedSetCommands interface { ZScan(key string, cursor string) (Result[string], []Result[string], error) ZScanWithOptions(key string, cursor string, options *options.ZScanOptions) (Result[string], []Result[string], error) + + ZRemRangeByLex(key string, rangeQuery options.RangeByLex) (int64, error) + + ZRemRangeByRank(key string, start int64, stop int64) (int64, error) + + ZRemRangeByScore(key string, rangeQuery options.RangeByScore) (int64, error) } diff --git a/go/integTest/shared_commands_test.go b/go/integTest/shared_commands_test.go index 6f3d019872..32dcb6a055 100644 --- a/go/integTest/shared_commands_test.go +++ b/go/integTest/shared_commands_test.go @@ -5976,3 +5976,162 @@ func (suite *GlideTestSuite) TestEcho() { assert.Equal(t, value, resultEcho.Value()) }) } + +func (suite *GlideTestSuite) TestZRemRangeByRank() { + suite.runWithDefaultClients(func(client api.BaseClient) { + key1 := uuid.New().String() + stringKey := uuid.New().String() + membersScores := map[string]float64{ + "one": 1.0, + "two": 2.0, + "three": 3.0, + } + zAddResult, err := client.ZAdd(key1, membersScores) + assert.NoError(suite.T(), err) + assert.Equal(suite.T(), int64(3), zAddResult) + + // Incorrect range start > stop + zRemRangeByRankResult, err := client.ZRemRangeByRank(key1, 2, 1) + assert.NoError(suite.T(), err) + assert.Equal(suite.T(), int64(0), zRemRangeByRankResult) + + // Remove first two members + zRemRangeByRankResult, err = client.ZRemRangeByRank(key1, 0, 1) + assert.NoError(suite.T(), err) + assert.Equal(suite.T(), int64(2), zRemRangeByRankResult) + + // Remove all members + zRemRangeByRankResult, err = client.ZRemRangeByRank(key1, 0, 10) + assert.NoError(suite.T(), err) + assert.Equal(suite.T(), int64(1), zRemRangeByRankResult) + + zRangeWithScoresResult, err := client.ZRangeWithScores(key1, options.NewRangeByIndexQuery(0, -1)) + assert.NoError(suite.T(), err) + assert.Equal(suite.T(), 0, len(zRangeWithScoresResult)) + + // Non-existing key + zRemRangeByRankResult, err = client.ZRemRangeByRank("non_existing_key", 0, 10) + assert.NoError(suite.T(), err) + assert.Equal(suite.T(), int64(0), zRemRangeByRankResult) + + // Key exists, but it is not a set + setResult, err := client.Set(stringKey, "test") + assert.NoError(suite.T(), err) + assert.Equal(suite.T(), "OK", setResult) + + _, err = client.ZRemRangeByRank(stringKey, 0, 10) + assert.NotNil(suite.T(), err) + assert.IsType(suite.T(), &api.RequestError{}, err) + }) +} + +func (suite *GlideTestSuite) TestZRemRangeByLex() { + suite.runWithDefaultClients(func(client api.BaseClient) { + key1 := uuid.New().String() + stringKey := uuid.New().String() + + // Add members to the set + zAddResult, err := client.ZAdd(key1, map[string]float64{"a": 1.0, "b": 2.0, "c": 3.0, "d": 4.0}) + assert.NoError(suite.T(), err) + assert.Equal(suite.T(), int64(4), zAddResult) + + // min > max + zRemRangeByLexResult, err := client.ZRemRangeByLex( + key1, + *options.NewRangeByLexQuery(options.NewLexBoundary("d", false), options.NewLexBoundary("a", false)), + ) + assert.NoError(suite.T(), err) + assert.Equal(suite.T(), int64(0), zRemRangeByLexResult) + + // Remove members with lexicographical range + zRemRangeByLexResult, err = client.ZRemRangeByLex( + key1, + *options.NewRangeByLexQuery(options.NewLexBoundary("a", false), options.NewLexBoundary("c", true)), + ) + assert.NoError(suite.T(), err) + assert.Equal(suite.T(), int64(2), zRemRangeByLexResult) + + zRemRangeByLexResult, err = client.ZRemRangeByLex( + key1, + *options.NewRangeByLexQuery(options.NewLexBoundary("d", true), options.NewInfiniteLexBoundary(options.PositiveInfinity)), + ) + assert.NoError(suite.T(), err) + assert.Equal(suite.T(), int64(1), zRemRangeByLexResult) + + // Non-existing key + zRemRangeByLexResult, err = client.ZRemRangeByLex( + "non_existing_key", + *options.NewRangeByLexQuery(options.NewLexBoundary("a", false), options.NewLexBoundary("c", false)), + ) + assert.NoError(suite.T(), err) + assert.Equal(suite.T(), int64(0), zRemRangeByLexResult) + + // Key exists, but it is not a set + setResult, err := client.Set(stringKey, "test") + assert.NoError(suite.T(), err) + assert.Equal(suite.T(), "OK", setResult) + + _, err = client.ZRemRangeByLex( + stringKey, + *options.NewRangeByLexQuery(options.NewLexBoundary("a", false), options.NewLexBoundary("c", false)), + ) + assert.NotNil(suite.T(), err) + assert.IsType(suite.T(), &api.RequestError{}, err) + }) +} + +func (suite *GlideTestSuite) TestZRemRangeByScore() { + suite.runWithDefaultClients(func(client api.BaseClient) { + key1 := uuid.New().String() + stringKey := uuid.New().String() + + // Add members to the set + zAddResult, err := client.ZAdd(key1, map[string]float64{"one": 1.0, "two": 2.0, "three": 3.0, "four": 4.0}) + assert.NoError(suite.T(), err) + assert.Equal(suite.T(), int64(4), zAddResult) + + // min > max + zRemRangeByScoreResult, err := client.ZRemRangeByScore( + key1, + *options.NewRangeByScoreQuery(options.NewScoreBoundary(2.0, false), options.NewScoreBoundary(1.0, false)), + ) + assert.NoError(suite.T(), err) + assert.Equal(suite.T(), int64(0), zRemRangeByScoreResult) + + // Remove members with score range + zRemRangeByScoreResult, err = client.ZRemRangeByScore( + key1, + *options.NewRangeByScoreQuery(options.NewScoreBoundary(1.0, false), options.NewScoreBoundary(3.0, true)), + ) + assert.NoError(suite.T(), err) + assert.Equal(suite.T(), int64(2), zRemRangeByScoreResult) + + // Remove all members + zRemRangeByScoreResult, err = client.ZRemRangeByScore( + key1, + *options.NewRangeByScoreQuery(options.NewScoreBoundary(1.0, false), options.NewScoreBoundary(10.0, true)), + ) + assert.NoError(suite.T(), err) + assert.Equal(suite.T(), int64(1), zRemRangeByScoreResult) + + // Non-existing key + zRemRangeByScoreResult, err = client.ZRemRangeByScore( + "non_existing_key", + *options.NewRangeByScoreQuery(options.NewScoreBoundary(1.0, false), options.NewScoreBoundary(10.0, true)), + ) + assert.NoError(suite.T(), err) + assert.Equal(suite.T(), int64(0), zRemRangeByScoreResult) + + // Key exists, but it is not a set + setResult, err := client.Set(stringKey, "test") + assert.NoError(suite.T(), err) + assert.Equal(suite.T(), "OK", setResult) + + _, err = client.ZRemRangeByScore( + stringKey, + *options.NewRangeByScoreQuery(options.NewScoreBoundary(1.0, false), options.NewScoreBoundary(10.0, true)), + ) + assert.NotNil(suite.T(), err) + assert.IsType(suite.T(), &api.RequestError{}, err) + }) +}