Java: Add Zrandmember command. (Sorted Set Commands) (valkey-io#1238)

* Java: Add Zrandmember command. (Sorted Set Commands) (#175) Signed-off-by: Yury-Fridlyand <[email protected]> Co-authored-by: Yury-Fridlyand <[email protected]> Co-authored-by: Aaron <[email protected]>
Bit-Quill · May 10, 2024 · d4ffa99 · d4ffa99
1 parent 0e1708a
commit d4ffa99
Show file tree

Hide file tree

Showing 10 changed files with 355 additions and 2 deletions.
diff --git a/glide-core/src/client/value_conversion.rs b/glide-core/src/client/value_conversion.rs
@@ -408,6 +408,9 @@ pub(crate) fn expected_type_for_cmd(cmd: &Cmd) -> Option<ExpectedReturnType> {
         b"HRANDFIELD" => cmd
             .position(b"WITHVALUES")
             .map(|_| ExpectedReturnType::ArrayOfKeyValuePairs),
+        b"ZRANDMEMBER" => cmd
+            .position(b"WITHSCORES")
+            .map(|_| ExpectedReturnType::ArrayOfKeyValuePairs),
         b"ZADD" => cmd
             .position(b"INCR")
             .map(|_| ExpectedReturnType::DoubleOrNull),
@@ -514,7 +517,7 @@ mod tests {
     }
 
     #[test]
-    fn convert_hrandfield() {
+    fn convert_array_of_kv_pairs() {
         assert!(matches!(
             expected_type_for_cmd(
                 redis::cmd("HRANDFIELD")
@@ -528,6 +531,19 @@ mod tests {
         assert!(expected_type_for_cmd(redis::cmd("HRANDFIELD").arg("key").arg("1")).is_none());
         assert!(expected_type_for_cmd(redis::cmd("HRANDFIELD").arg("key")).is_none());
 
+        assert!(matches!(
+            expected_type_for_cmd(
+                redis::cmd("ZRANDMEMBER")
+                    .arg("key")
+                    .arg("1")
+                    .arg("withscores")
+            ),
+            Some(ExpectedReturnType::ArrayOfKeyValuePairs)
+        ));
+
+        assert!(expected_type_for_cmd(redis::cmd("ZRANDMEMBER").arg("key").arg("1")).is_none());
+        assert!(expected_type_for_cmd(redis::cmd("ZRANDMEMBER").arg("key")).is_none());
+
         let flat_array = Value::Array(vec![
             Value::BulkString(b"key1".to_vec()),
             Value::BulkString(b"value1".to_vec()),

diff --git a/glide-core/src/protobuf/redis_request.proto b/glide-core/src/protobuf/redis_request.proto
@@ -180,6 +180,7 @@ enum RequestType {
     ZUnion = 136;
     BZPopMin = 137;
     FlushAll = 138;
+    ZRandMember = 139;
 }
 
 message Command {

diff --git a/glide-core/src/request_type.rs b/glide-core/src/request_type.rs
@@ -148,6 +148,7 @@ pub enum RequestType {
     ZUnion = 136,
     BZPopMin = 137,
     FlushAll = 138,
+    ZRandMember = 139,
 }
 
 fn get_two_word_command(first: &str, second: &str) -> Cmd {
@@ -299,6 +300,7 @@ impl From<::protobuf::EnumOrUnknown<ProtobufRequestType>> for RequestType {
             ProtobufRequestType::ZUnion => RequestType::ZUnion,
             ProtobufRequestType::BZPopMin => RequestType::BZPopMin,
             ProtobufRequestType::FlushAll => RequestType::FlushAll,
+            ProtobufRequestType::ZRandMember => RequestType::ZRandMember,
         }
     }
 }
@@ -446,6 +448,7 @@ impl RequestType {
             RequestType::ZUnion => Some(cmd("ZUNION")),
             RequestType::BZPopMin => Some(cmd("BZPOPMIN")),
             RequestType::FlushAll => Some(cmd("FLUSHALL")),
+            RequestType::ZRandMember => Some(cmd("ZRANDMEMBER")),
         }
     }
 }
diff --git a/java/client/src/main/java/glide/api/BaseClient.java b/java/client/src/main/java/glide/api/BaseClient.java
@@ -91,6 +91,7 @@
 import static redis_request.RedisRequestOuterClass.RequestType.ZMScore;
 import static redis_request.RedisRequestOuterClass.RequestType.ZPopMax;
 import static redis_request.RedisRequestOuterClass.RequestType.ZPopMin;
+import static redis_request.RedisRequestOuterClass.RequestType.ZRandMember;
 import static redis_request.RedisRequestOuterClass.RequestType.ZRangeStore;
 import static redis_request.RedisRequestOuterClass.RequestType.ZRemRangeByLex;
 import static redis_request.RedisRequestOuterClass.RequestType.ZRemRangeByRank;
@@ -1037,6 +1038,30 @@ public CompletableFuture<Map<String, Double>> zunionWithScores(
         return commandManager.submitNewCommand(ZUnion, arguments, this::handleMapResponse);
     }
 
+    @Override
+    public CompletableFuture<String> zrandmember(@NonNull String key) {
+        return commandManager.submitNewCommand(
+                ZRandMember, new String[] {key}, this::handleStringOrNullResponse);
+    }
+
+    @Override
+    public CompletableFuture<String[]> zrandmemberWithCount(@NonNull String key, long count) {
+        return commandManager.submitNewCommand(
+                ZRandMember,
+                new String[] {key, Long.toString(count)},
+                response -> castArray(handleArrayResponse(response), String.class));
+    }
+
+    @Override
+    public CompletableFuture<Object[][]> zrandmemberWithCountWithScores(
+            @NonNull String key, long count) {
+        String[] arguments = new String[] {key, Long.toString(count), WITH_SCORES_REDIS_API};
+        return commandManager.submitNewCommand(
+                ZRandMember,
+                arguments,
+                response -> castArray(handleArrayResponse(response), Object[].class));
+    }
+
     @Override
     public CompletableFuture<String> xadd(@NonNull String key, @NonNull Map<String, String> values) {
         return xadd(key, values, StreamAddOptions.builder().build());

diff --git a/java/client/src/main/java/glide/api/commands/SortedSetBaseCommands.java b/java/client/src/main/java/glide/api/commands/SortedSetBaseCommands.java
@@ -1065,4 +1065,68 @@ CompletableFuture<Map<String, Double>> zunionWithScores(
      * }</pre>
      */
     CompletableFuture<Map<String, Double>> zunionWithScores(KeysOrWeightedKeys keysOrWeightedKeys);
+
+    /**
+     * Returns a random element from the sorted set stored at <code>key</code>.
+     *
+     * @see <a href="https://redis.io/commands/zrandmember/">redis.io</a> for more details.
+     * @param key The key of the sorted set.
+     * @return A <code>String</code> representing a random element from the sorted set.<br>
+     *     If the sorted set does not exist or is empty, the response will be <code>null</code>.
+     * @example
+     *     <pre>{@code
+     * String payload1 = client.zrandmember("mySortedSet").get();
+     * assert payload1.equals("GLIDE");
+     *
+     * String payload2 = client.zrandmember("nonExistingSortedSet").get();
+     * assert payload2 == null;
+     * }</pre>
+     */
+    CompletableFuture<String> zrandmember(String key);
+
+    /**
+     * Retrieves random elements from the sorted set stored at <code>key</code>.
+     *
+     * @see <a href="https://redis.io/commands/zrandmember/">redis.io</a> for more details.
+     * @param key The key of the sorted set.
+     * @param count The number of elements to return.<br>
+     *     If <code>count</code> is positive, returns unique elements.<br>
+     *     If negative, allows for duplicates.<br>
+     * @return An <code>array</code> of elements from the sorted set.<br>
+     *     If the sorted set does not exist or is empty, the response will be an empty <code>array
+     *     </code>.
+     * @example
+     *     <pre>{@code
+     * String[] payload1 = client.zrandmember("mySortedSet", -3).get();
+     * assert payload1.equals(new String[] {"GLIDE", "GLIDE", "JAVA"});
+     *
+     * String[] payload2 = client.zrandmember("nonExistingSortedSet", 3).get();
+     * assert payload2.length == 0;
+     * }</pre>
+     */
+    CompletableFuture<String[]> zrandmemberWithCount(String key, long count);
+
+    /**
+     * Retrieves random elements along with their scores from the sorted set stored at <code>key
+     * </code>.
+     *
+     * @see <a href="https://redis.io/commands/zrandmember/">redis.io</a> for more details.
+     * @param key The key of the sorted set.
+     * @param count The number of elements to return.<br>
+     *     If <code>count</code> is positive, returns unique elements.<br>
+     *     If negative, allows duplicates.<br>
+     * @return An <code>array</code> of <code>[element, score]</code> <code>arrays</code>, where
+     *     element is a <code>String</code> and score is a <code>Double</code>.<br>
+     *     If the sorted set does not exist or is empty, the response will be an empty <code>array
+     *     </code>.
+     * @example
+     *     <pre>{@code
+     * Object[][] data = client.zrandmemberWithCountWithScores(key1, -3).get();
+     * assert data.length == 3;
+     * for (Object[] memberScorePair : data) {
+     *     System.out.printf("Member: '%s', score: %d", memberScorePair[0], memberScorePair[1]);
+     * }
+     * }</pre>
+     */
+    CompletableFuture<Object[][]> zrandmemberWithCountWithScores(String key, long count);
 }
diff --git a/java/client/src/main/java/glide/api/models/BaseTransaction.java b/java/client/src/main/java/glide/api/models/BaseTransaction.java
@@ -106,6 +106,7 @@
 import static redis_request.RedisRequestOuterClass.RequestType.ZMScore;
 import static redis_request.RedisRequestOuterClass.RequestType.ZPopMax;
 import static redis_request.RedisRequestOuterClass.RequestType.ZPopMin;
+import static redis_request.RedisRequestOuterClass.RequestType.ZRandMember;
 import static redis_request.RedisRequestOuterClass.RequestType.ZRangeStore;
 import static redis_request.RedisRequestOuterClass.RequestType.ZRemRangeByLex;
 import static redis_request.RedisRequestOuterClass.RequestType.ZRemRangeByRank;
@@ -1601,6 +1602,61 @@ public T zpopmin(@NonNull String key) {
         return getThis();
     }
 
+    /**
+     * Returns a random element from the sorted set stored at <code>key</code>.
+     *
+     * @see <a href="https://redis.io/commands/zrandmember/">redis.io</a> for more details.
+     * @param key The key of the sorted set.
+     * @return Command Response - A <code>String</code> representing a random element from the sorted
+     *     set.<br>
+     *     If the sorted set does not exist or is empty, the response will be <code>null</code>.
+     */
+    public T zrandmember(@NonNull String key) {
+        ArgsArray commandArgs = buildArgs(key);
+        protobufTransaction.addCommands(buildCommand(ZRandMember, commandArgs));
+        return getThis();
+    }
+
+    /**
+     * Retrieves random elements from the sorted set stored at <code>key</code>.
+     *
+     * @see <a href="https://redis.io/commands/zrandmember/">redis.io</a> for more details.
+     * @param key The key of the sorted set.
+     * @param count The number of elements to return.<br>
+     *     If <code>count</code> is positive, returns unique elements.<br>
+     *     If negative, allows for duplicates.<br>
+     * @return Command Response - An <code>array</code> of elements from the sorted set.<br>
+     *     If the sorted set does not exist or is empty, the response will be an empty <code>array
+     *     </code>.
+     */
+    public T zrandmemberWithCount(@NonNull String key, long count) {
+        ArgsArray commandArgs = buildArgs(key, Long.toString(count));
+        protobufTransaction.addCommands(buildCommand(ZRandMember, commandArgs));
+        return getThis();
+    }
+
+    /**
+     * Retrieves random elements along with their scores from the sorted set stored at <code>key
+     * </code>.
+     *
+     * @see <a href="https://redis.io/commands/zrandmember/">redis.io</a> for more details.
+     * @param key The key of the sorted set.
+     * @param count The number of elements to return.<br>
+     *     If <code>count</code> is positive, returns unique elements.<br>
+     *     If negative, allows duplicates.<br>
+     * @return Command Response - An <code>array</code> of <code>[element, score]</code> <code>arrays
+     *     </code>, where element is a <code>String</code> and score is a <code>Double</code>.<br>
+     *     If the sorted set does not exist or is empty, the response will be an empty <code>array
+     *     </code>.
+     */
+    public T zrandmemberWithCountWithScores(String key, long count) {
+        String[] arguments = new String[] {key, Long.toString(count), WITH_SCORES_REDIS_API};
+
+        ArgsArray commandArgs = buildArgs(arguments);
+        protobufTransaction.addCommands(buildCommand(ZRandMember, commandArgs));
+        return getThis();
+    }
+
     /**
      * Blocks the connection until it removes and returns a member with the lowest score from the
      * sorted sets stored at the specified <code>keys</code>. The sorted sets are checked in the order

diff --git a/java/client/src/test/java/glide/api/RedisClientTest.java b/java/client/src/test/java/glide/api/RedisClientTest.java
@@ -128,6 +128,7 @@
 import static redis_request.RedisRequestOuterClass.RequestType.ZMScore;
 import static redis_request.RedisRequestOuterClass.RequestType.ZPopMax;
 import static redis_request.RedisRequestOuterClass.RequestType.ZPopMin;
+import static redis_request.RedisRequestOuterClass.RequestType.ZRandMember;
 import static redis_request.RedisRequestOuterClass.RequestType.ZRangeStore;
 import static redis_request.RedisRequestOuterClass.RequestType.ZRemRangeByLex;
 import static redis_request.RedisRequestOuterClass.RequestType.ZRemRangeByRank;
@@ -3252,11 +3253,84 @@ public void xadd_returns_success() {
 
         // exercise
         CompletableFuture<String> response = service.xadd(key, fieldValues);
+
+        // verify
+        assertEquals(testResponse, response);
+        assertEquals(returnId, response.get());
+    }
+
+    @SneakyThrows
+    @Test
+    public void zrandmember_returns_success() {
+        // setup
+        String key = "testKey";
+        String[] arguments = new String[] {key};
+        String value = "testValue";
+
+        CompletableFuture<String> testResponse = new CompletableFuture<>();
+        testResponse.complete(value);
+
+        // match on protobuf request
+        when(commandManager.<String>submitNewCommand(eq(ZRandMember), eq(arguments), any()))
+                .thenReturn(testResponse);
+
+        // exercise
+        CompletableFuture<String> response = service.zrandmember(key);
         String payload = response.get();
 
         // verify
         assertEquals(testResponse, response);
-        assertEquals(returnId, payload);
+        assertEquals(value, payload);
+    }
+
+    @SneakyThrows
+    @Test
+    public void zrandmemberWithCount_returns_success() {
+        // setup
+        String key = "testKey";
+        long count = 2L;
+        String[] arguments = new String[] {key, Long.toString(count)};
+        String[] value = new String[] {"member1", "member2"};
+
+        CompletableFuture<String[]> testResponse = new CompletableFuture<>();
+        testResponse.complete(value);
+
+        // match on protobuf request
+        when(commandManager.<String[]>submitNewCommand(eq(ZRandMember), eq(arguments), any()))
+                .thenReturn(testResponse);
+
+        // exercise
+        CompletableFuture<String[]> response = service.zrandmemberWithCount(key, count);
+        String[] payload = response.get();
+
+        // verify
+        assertEquals(testResponse, response);
+        assertEquals(value, payload);
+    }
+
+    @SneakyThrows
+    @Test
+    public void zrandmemberWithCountWithScores_returns_success() {
+        // setup
+        String key = "testKey";
+        long count = 2L;
+        String[] arguments = new String[] {key, Long.toString(count), WITH_SCORES_REDIS_API};
+        Object[][] value = new Object[][] {{"member1", 2.0}, {"member2", 3.0}};
+
+        CompletableFuture<Object[][]> testResponse = new CompletableFuture<>();
+        testResponse.complete(value);
+
+        // match on protobuf request
+        when(commandManager.<Object[][]>submitNewCommand(eq(ZRandMember), eq(arguments), any()))
+                .thenReturn(testResponse);
+
+        // exercise
+        CompletableFuture<Object[][]> response = service.zrandmemberWithCountWithScores(key, count);
+        Object[] payload = response.get();
+
+        // verify
+        assertEquals(testResponse, response);
+        assertEquals(value, payload);
     }
 
     private static List<Arguments> getStreamAddOptions() {

diff --git a/java/client/src/test/java/glide/api/models/TransactionTests.java b/java/client/src/test/java/glide/api/models/TransactionTests.java
@@ -116,6 +116,7 @@
 import static redis_request.RedisRequestOuterClass.RequestType.ZMScore;
 import static redis_request.RedisRequestOuterClass.RequestType.ZPopMax;
 import static redis_request.RedisRequestOuterClass.RequestType.ZPopMin;
+import static redis_request.RedisRequestOuterClass.RequestType.ZRandMember;
 import static redis_request.RedisRequestOuterClass.RequestType.ZRangeStore;
 import static redis_request.RedisRequestOuterClass.RequestType.ZRemRangeByLex;
 import static redis_request.RedisRequestOuterClass.RequestType.ZRemRangeByRank;
@@ -595,6 +596,22 @@ InfScoreBound.NEGATIVE_INFINITY, new ScoreBoundary(3, false), new Limit(1, 2)),
         transaction.persist("key");
         results.add(Pair.of(Persist, buildArgs("key")));
 
+        transaction.zrandmember("key");
+        results.add(Pair.of(ZRandMember, ArgsArray.newBuilder().addArgs("key").build()));
+
+        transaction.zrandmemberWithCount("key", 5);
+        results.add(Pair.of(ZRandMember, ArgsArray.newBuilder().addArgs("key").addArgs("5").build()));
+
+        transaction.zrandmemberWithCountWithScores("key", 5);
+        results.add(
+                Pair.of(
+                        ZRandMember,
+                        ArgsArray.newBuilder()
+                                .addArgs("key")
+                                .addArgs("5")
+                                .addArgs(WITH_SCORES_REDIS_API)
+                                .build()));
+
         transaction.type("key");
         results.add(Pair.of(Type, buildArgs("key")));