Go: Add commands ZRemRangeByRank/ZRemRangeByScore/ZRemRangeByLex (#2967)

Signed-off-by: TJ Zhang <[email protected]> Co-authored-by: TJ Zhang <[email protected]>
valkey-io · Jan 17, 2025 · 4bf3cf4 · 4bf3cf4
1 parent 65667f4
commit 4bf3cf4
Show file tree

Hide file tree

Showing 4 changed files with 272 additions and 0 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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)
+	})
+}