Java: add PUBSUB SHARDNUMSUB (valkey-io#2279)

* add PUBSUB SHARDNUMSUB in Java --------- Signed-off-by: James Xin <james.xin@improving.com>
shohamazon · Sep 12, 2024 · 8bf8282 · 8bf8282
1 parent e997eea
commit 8bf8282
Show file tree

Hide file tree

Showing 7 changed files with 174 additions and 0 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -26,6 +26,7 @@
 * Node: Added XINFO GROUPS command ([#2122](https://github.com/valkey-io/valkey-glide/pull/2122))
 * Java: Added PUBSUB CHANNELS, NUMPAT and NUMSUB commands ([#2105](https://github.com/valkey-io/valkey-glide/pull/2105))
 * Java: Added PUBSUB SHARDCHANNELS command ([#2265](https://github.com/valkey-io/valkey-glide/pull/2265))
+* Java: Added PUBSUB SHARDNUMSUB command ([#2279](https://github.com/valkey-io/valkey-glide/pull/2279))
 * Java: Added binary support for custom command ([#2109](https://github.com/valkey-io/valkey-glide/pull/2109))
 * Node: Added SSCAN command ([#2132](https://github.com/valkey-io/valkey-glide/pull/2132))
 * Node: Added HKEYS command ([#2136](https://github.com/valkey-io/valkey-glide/pull/2136))

diff --git a/java/client/src/main/java/glide/api/GlideClusterClient.java b/java/client/src/main/java/glide/api/GlideClusterClient.java
@@ -27,6 +27,7 @@
 import static command_request.CommandRequestOuterClass.RequestType.Lolwut;
 import static command_request.CommandRequestOuterClass.RequestType.Ping;
 import static command_request.CommandRequestOuterClass.RequestType.PubSubSChannels;
+import static command_request.CommandRequestOuterClass.RequestType.PubSubSNumSub;
 import static command_request.CommandRequestOuterClass.RequestType.RandomKey;
 import static command_request.CommandRequestOuterClass.RequestType.SPublish;
 import static command_request.CommandRequestOuterClass.RequestType.Sort;
@@ -1050,6 +1051,18 @@ public CompletableFuture<GlideString[]> pubsubShardChannels(@NonNull GlideString
                 response -> castArray(handleArrayResponseBinary(response), GlideString.class));
     }
 
+    @Override
+    public CompletableFuture<Map<String, Long>> pubsubShardNumSub(@NonNull String[] channels) {
+        return commandManager.submitNewCommand(PubSubSNumSub, channels, this::handleMapResponse);
+    }
+
+    @Override
+    public CompletableFuture<Map<GlideString, Long>> pubsubShardNumSub(
+            @NonNull GlideString[] channels) {
+        return commandManager.submitNewCommand(
+                PubSubSNumSub, channels, this::handleBinaryStringMapResponse);
+    }
+
     @Override
     public CompletableFuture<String> unwatch(@NonNull Route route) {
         return commandManager.submitNewCommand(

diff --git a/java/client/src/main/java/glide/api/commands/PubSubClusterCommands.java b/java/client/src/main/java/glide/api/commands/PubSubClusterCommands.java
@@ -2,6 +2,7 @@
 package glide.api.commands;
 
 import glide.api.models.GlideString;
+import java.util.Map;
 import java.util.concurrent.CompletableFuture;
 
 /**
@@ -100,4 +101,38 @@ public interface PubSubClusterCommands {
      * }</pre>
      */
     CompletableFuture<GlideString[]> pubsubShardChannels(GlideString pattern);
+
+    /**
+     * Returns the number of subscribers (exclusive of clients subscribed to patterns) for the
+     * specified shard channels. Note that it is valid to call this command without channels. In this
+     * case, it will just return an empty map.
+     *
+     * @see <a href="https://valkey.io/commands/pubsub-shardnumsub/">valkey.io</a> for details.
+     * @param channels The list of shard channels to query for the number of subscribers.
+     * @return An <code>Map</code> where keys are the shard channel names and values are the number of
+     *     subscribers.
+     * @example
+     *     <pre>{@code
+     * Map<String, Long> result = client.pubsubShardNumSub(new String[] {"channel1", "channel2"}).get();
+     * assert result.equals(Map.of("channel1", 3L, "channel2", 5L));
+     * }</pre>
+     */
+    CompletableFuture<Map<String, Long>> pubsubShardNumSub(String[] channels);
+
+    /**
+     * Returns the number of subscribers (exclusive of clients subscribed to patterns) for the
+     * specified shard channels. Note that it is valid to call this command without channels. In this
+     * case, it will just return an empty map.
+     *
+     * @see <a href="https://valkey.io/commands/pubsub-shardnumsub/">valkey.io</a> for details.
+     * @param channels The list of shard channels to query for the number of subscribers.
+     * @return An <code>Map</code> where keys are the shard channel names and values are the number of
+     *     subscribers.
+     * @example
+     *     <pre>{@code
+     * Map<GlideString, Long> result = client.pubsubShardNumSub(new GlideString[] {gs("channel1"), gs("channel2")}).get();
+     * assert result.equals(Map.of(gs("channel1"), 3L, gs("channel2"), 5L));
+     * }</pre>
+     */
+    CompletableFuture<Map<GlideString, Long>> pubsubShardNumSub(GlideString[] channels);
 }
diff --git a/java/client/src/main/java/glide/api/models/ClusterTransaction.java b/java/client/src/main/java/glide/api/models/ClusterTransaction.java
@@ -2,6 +2,7 @@
 package glide.api.models;
 
 import static command_request.CommandRequestOuterClass.RequestType.PubSubSChannels;
+import static command_request.CommandRequestOuterClass.RequestType.PubSubSNumSub;
 import static command_request.CommandRequestOuterClass.RequestType.SPublish;
 import static command_request.CommandRequestOuterClass.RequestType.Sort;
 import static command_request.CommandRequestOuterClass.RequestType.SortReadOnly;
@@ -60,6 +61,24 @@ public <ArgType> ClusterTransaction publish(
         return getThis();
     }
 
+    /**
+     * Returns the number of subscribers (exclusive of clients subscribed to patterns) for the
+     * specified shard channels. Note that it is valid to call this command without channels. In this
+     * case, it will just return an empty map.
+     *
+     * @implNote {@link ArgType} is limited to {@link String} or {@link GlideString}, any other type
+     *     will throw {@link IllegalArgumentException}.
+     * @see <a href="https://valkey.io/commands/pubsub-shardnumsub/">valkey.io</a> for details.
+     * @param channels The list of shard channels to query for the number of subscribers.
+     * @return Command response - An <code>Map</code> where keys are the shard channel names and
+     *     values are the number of subscribers.
+     */
+    public <ArgType> ClusterTransaction pubsubShardNumSub(@NonNull ArgType[] channels) {
+        checkTypeOrThrow(channels);
+        protobufTransaction.addCommands(buildCommand(PubSubSNumSub, newArgsBuilder().add(channels)));
+        return getThis();
+    }
+
     /**
      * Sorts the elements in the list, set, or sorted set at <code>key</code> and returns the result.
      * <br>

diff --git a/java/client/src/test/java/glide/api/GlideClusterClientTest.java b/java/client/src/test/java/glide/api/GlideClusterClientTest.java
@@ -26,6 +26,7 @@
 import static command_request.CommandRequestOuterClass.RequestType.Lolwut;
 import static command_request.CommandRequestOuterClass.RequestType.Ping;
 import static command_request.CommandRequestOuterClass.RequestType.PubSubSChannels;
+import static command_request.CommandRequestOuterClass.RequestType.PubSubSNumSub;
 import static command_request.CommandRequestOuterClass.RequestType.RandomKey;
 import static command_request.CommandRequestOuterClass.RequestType.SPublish;
 import static command_request.CommandRequestOuterClass.RequestType.Sort;
@@ -2812,6 +2813,54 @@ public void pubsubShardChannelsBinary_with_pattern_returns_success() {
         assertEquals(value, payload);
     }
 
+    @SneakyThrows
+    @Test
+    public void pubsubShardNumSub_returns_success() {
+        // setup
+        String[] arguments = new String[] {"ch1", "ch2"};
+        Map<String, Long> value = Map.of();
+
+        CompletableFuture<Map<String, Long>> testResponse = new CompletableFuture<>();
+        testResponse.complete(value);
+
+        // match on protobuf request
+        when(commandManager.<Map<String, Long>>submitNewCommand(
+                        eq(PubSubSNumSub), eq(arguments), any()))
+                .thenReturn(testResponse);
+
+        // exercise
+        CompletableFuture<Map<String, Long>> response = service.pubsubShardNumSub(arguments);
+        Map<String, Long> payload = response.get();
+
+        // verify
+        assertEquals(testResponse, response);
+        assertEquals(value, payload);
+    }
+
+    @SneakyThrows
+    @Test
+    public void pubsubShardNumSub_binary_returns_success() {
+        // setup
+        GlideString[] arguments = new GlideString[] {gs("ch1"), gs("ch2")};
+        Map<GlideString, Long> value = Map.of();
+
+        CompletableFuture<Map<GlideString, Long>> testResponse = new CompletableFuture<>();
+        testResponse.complete(value);
+
+        // match on protobuf request
+        when(commandManager.<Map<GlideString, Long>>submitNewCommand(
+                        eq(PubSubSNumSub), eq(arguments), any()))
+                .thenReturn(testResponse);
+
+        // exercise
+        CompletableFuture<Map<GlideString, Long>> response = service.pubsubShardNumSub(arguments);
+        Map<GlideString, Long> payload = response.get();
+
+        // verify
+        assertEquals(testResponse, response);
+        assertEquals(value, payload);
+    }
+
     @SneakyThrows
     @Test
     public void sort_returns_success() {

diff --git a/java/client/src/test/java/glide/api/models/ClusterTransactionTests.java b/java/client/src/test/java/glide/api/models/ClusterTransactionTests.java
@@ -2,6 +2,7 @@
 package glide.api.models;
 
 import static command_request.CommandRequestOuterClass.RequestType.PubSubSChannels;
+import static command_request.CommandRequestOuterClass.RequestType.PubSubSNumSub;
 import static command_request.CommandRequestOuterClass.RequestType.SPublish;
 import static command_request.CommandRequestOuterClass.RequestType.Sort;
 import static command_request.CommandRequestOuterClass.RequestType.SortReadOnly;
@@ -38,6 +39,9 @@ public void cluster_transaction_builds_protobuf_request() {
         transaction.pubsubShardChannels("test*");
         results.add(Pair.of(PubSubSChannels, buildArgs("test*")));
 
+        transaction.pubsubShardNumSub(new String[] {"ch1", "ch2"});
+        results.add(Pair.of(PubSubSNumSub, buildArgs("ch1", "ch2")));
+
         transaction.sortReadOnly(
                 "key1",
                 SortClusterOptions.builder()

diff --git a/java/integTest/src/test/java/glide/PubSubTests.java b/java/integTest/src/test/java/glide/PubSubTests.java
@@ -1523,4 +1523,57 @@ public void pubsub_shard_channels() {
         assertEquals(0, listener.pubsubShardChannels("non_matching_*").get().length);
         assertEquals(0, listener.pubsubShardChannels(gs("non_matching_*")).get().length);
     }
+
+    @SneakyThrows
+    @Test
+    public void pubsub_shardnumsub() {
+        assumeTrue(SERVER_VERSION.isGreaterThanOrEqualTo("7.0.0"), "This feature added in version 7");
+
+        // no channels exists yet
+        GlideClusterClient client = (GlideClusterClient) createClient(false);
+        var channels = new String[] {"channel1", "channel2", "channel3"};
+        assertEquals(
+                Arrays.stream(channels).collect(Collectors.toMap(c -> c, c -> 0L)),
+                client.pubsubNumSub(channels).get());
+
+        Map<? extends ChannelMode, Set<GlideString>> subscriptions1 =
+                Map.of(
+                        PubSubClusterChannelMode.SHARDED,
+                        Set.of(gs("channel1"), gs("channel2"), gs("channel3")));
+        GlideClusterClient listener1 =
+                (GlideClusterClient) createClientWithSubscriptions(false, subscriptions1);
+
+        Map<? extends ChannelMode, Set<GlideString>> subscriptions2 =
+                Map.of(PubSubClusterChannelMode.SHARDED, Set.of(gs("channel2"), gs("channel3")));
+        GlideClusterClient listener2 =
+                (GlideClusterClient) createClientWithSubscriptions(false, subscriptions2);
+
+        Map<? extends ChannelMode, Set<GlideString>> subscriptions3 =
+                Map.of(PubSubClusterChannelMode.SHARDED, Set.of(gs("channel3")));
+        GlideClusterClient listener3 =
+                (GlideClusterClient) createClientWithSubscriptions(false, subscriptions3);
+
+        clients.addAll(List.of(client, listener1, listener2, listener3));
+
+        var expected = Map.of("channel1", 1L, "channel2", 2L, "channel3", 3L, "channel4", 0L);
+        assertEquals(
+                expected, client.pubsubShardNumSub(ArrayUtils.addFirst(channels, "channel4")).get());
+        assertEquals(
+                expected, listener1.pubsubShardNumSub(ArrayUtils.addFirst(channels, "channel4")).get());
+
+        var expectedGs =
+                Map.of(gs("channel1"), 1L, gs("channel2"), 2L, gs("channel3"), 3L, gs("channel4"), 0L);
+        assertEquals(
+                expectedGs,
+                client
+                        .pubsubShardNumSub(
+                                new GlideString[] {gs("channel1"), gs("channel2"), gs("channel3"), gs("channel4")})
+                        .get());
+        assertEquals(
+                expectedGs,
+                listener2
+                        .pubsubShardNumSub(
+                                new GlideString[] {gs("channel1"), gs("channel2"), gs("channel3"), gs("channel4")})
+                        .get());
+    }
 }