Go: Add command ZRank and ZRevRank (#2932)

* Go: Add command ZRank and ZRevRank

Signed-off-by: TJ Zhang <[email protected]>

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)
+}

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.
 )

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)
 

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)
 }

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)
+	})
+}