Java: add DBSIZE command (#1540)

* Java: add DBSIZE command (#346) * PR suggestions
valkey-io · Jun 7, 2024 · 5ec5b71 · 5ec5b71
1 parent e2d35f2
commit 5ec5b71
Show file tree

Hide file tree

Showing 11 changed files with 184 additions and 3 deletions.
diff --git a/java/client/src/main/java/glide/api/RedisClient.java b/java/client/src/main/java/glide/api/RedisClient.java
@@ -14,6 +14,7 @@
 import static redis_request.RedisRequestOuterClass.RequestType.ConfigRewrite;
 import static redis_request.RedisRequestOuterClass.RequestType.ConfigSet;
 import static redis_request.RedisRequestOuterClass.RequestType.CustomCommand;
+import static redis_request.RedisRequestOuterClass.RequestType.DBSize;
 import static redis_request.RedisRequestOuterClass.RequestType.Echo;
 import static redis_request.RedisRequestOuterClass.RequestType.FlushAll;
 import static redis_request.RedisRequestOuterClass.RequestType.FunctionDelete;
@@ -195,6 +196,11 @@ public CompletableFuture<String> lolwut(int version, int @NonNull [] parameters)
         return commandManager.submitNewCommand(Lolwut, arguments, this::handleStringResponse);
     }
 
+    @Override
+    public CompletableFuture<Long> dbsize() {
+        return commandManager.submitNewCommand(DBSize, new String[0], this::handleLongResponse);
+    }
+
     @Override
     public CompletableFuture<String> functionLoad(@NonNull String libraryCode, boolean replace) {
         String[] arguments =

diff --git a/java/client/src/main/java/glide/api/RedisClusterClient.java b/java/client/src/main/java/glide/api/RedisClusterClient.java
@@ -16,6 +16,7 @@
 import static redis_request.RedisRequestOuterClass.RequestType.ConfigRewrite;
 import static redis_request.RedisRequestOuterClass.RequestType.ConfigSet;
 import static redis_request.RedisRequestOuterClass.RequestType.CustomCommand;
+import static redis_request.RedisRequestOuterClass.RequestType.DBSize;
 import static redis_request.RedisRequestOuterClass.RequestType.Echo;
 import static redis_request.RedisRequestOuterClass.RequestType.FlushAll;
 import static redis_request.RedisRequestOuterClass.RequestType.FunctionDelete;
@@ -422,6 +423,16 @@ public CompletableFuture<ClusterValue<String>> lolwut(
                                 : ClusterValue.ofMultiValue(handleMapResponse(response)));
     }
 
+    @Override
+    public CompletableFuture<Long> dbsize() {
+        return commandManager.submitNewCommand(DBSize, new String[0], this::handleLongResponse);
+    }
+
+    @Override
+    public CompletableFuture<Long> dbsize(@NonNull Route route) {
+        return commandManager.submitNewCommand(DBSize, new String[0], route, this::handleLongResponse);
+    }
+
     @Override
     public CompletableFuture<String> functionLoad(@NonNull String libraryCode, boolean replace) {
         String[] arguments =

diff --git a/java/client/src/main/java/glide/api/commands/ServerManagementClusterCommands.java b/java/client/src/main/java/glide/api/commands/ServerManagementClusterCommands.java
@@ -560,4 +560,36 @@ public interface ServerManagementClusterCommands {
      * }</pre>
      */
     CompletableFuture<ClusterValue<String>> lolwut(int version, int[] parameters, Route route);
+
+    /**
+     * Returns the number of keys in the database.<br>
+     * The command will be routed to all primary nodes.
+     *
+     * @see <a href="https://valkey.io/commands/dbsize/">valkey.io</a> for details.
+     * @return The total number of keys across the primary nodes.
+     * @example
+     *     <pre>{@code
+     * Long numKeys = client.dbsize().get();
+     * System.out.printf("Number of keys across the primary nodes: %d%n", numKeys);
+     * }</pre>
+     */
+    CompletableFuture<Long> dbsize();
+
+    /**
+     * Returns the number of keys in the database.
+     *
+     * @see <a href="https://valkey.io/commands/dbsize/">valkey.io</a> for details.
+     * @param route Specifies the routing configuration for the command. The client will route the
+     *     command to the nodes defined by <code>route</code>.
+     * @return The number of keys in the database.<br>
+     *     If the query is routed to multiple nodes, returns the sum of the number of keys across all
+     *     routed nodes.
+     * @example
+     *     <pre>{@code
+     * Route route = new ByAddressRoute("localhost", 8000);
+     * Long numKeys = client.dbsize(route).get();
+     * System.out.printf("Number of keys for node at port 8000: %d%n", numKeys);
+     * }</pre>
+     */
+    CompletableFuture<Long> dbsize(Route route);
 }
diff --git a/java/client/src/main/java/glide/api/commands/ServerManagementCommands.java b/java/client/src/main/java/glide/api/commands/ServerManagementCommands.java
@@ -257,4 +257,17 @@ public interface ServerManagementCommands {
      * }</pre>
      */
     CompletableFuture<String> lolwut(int version, int[] parameters);
+
+    /**
+     * Returns the number of keys in the currently selected database.
+     *
+     * @see <a href="https://valkey.io/commands/dbsize/">valkey.io</a> for details.
+     * @return The number of keys in the currently selected database.
+     * @example
+     *     <pre>{@code
+     * Long numKeys = client.dbsize().get();
+     * System.out.printf("Number of keys in the current database: %d%n", numKeys);
+     * }</pre>
+     */
+    CompletableFuture<Long> dbsize();
 }
diff --git a/java/client/src/main/java/glide/api/models/BaseTransaction.java b/java/client/src/main/java/glide/api/models/BaseTransaction.java
@@ -38,6 +38,7 @@
 import static redis_request.RedisRequestOuterClass.RequestType.ConfigRewrite;
 import static redis_request.RedisRequestOuterClass.RequestType.ConfigSet;
 import static redis_request.RedisRequestOuterClass.RequestType.CustomCommand;
+import static redis_request.RedisRequestOuterClass.RequestType.DBSize;
 import static redis_request.RedisRequestOuterClass.RequestType.Decr;
 import static redis_request.RedisRequestOuterClass.RequestType.DecrBy;
 import static redis_request.RedisRequestOuterClass.RequestType.Del;
@@ -2956,6 +2957,17 @@ public T lolwut(int version, int @NonNull [] parameters) {
         return getThis();
     }
 
+    /**
+     * Returns the number of keys in the currently selected database.
+     *
+     * @see <a href="https://valkey.io/commands/dbsize/">valkey.io</a> for details.
+     * @return Command Response - The number of keys in the currently selected database.
+     */
+    public T dbsize() {
+        protobufTransaction.addCommands(buildCommand(DBSize));
+        return getThis();
+    }
+
     /**
      * Returns the string representation of the type of the value stored at <code>key</code>.
      *

diff --git a/java/client/src/test/java/glide/api/RedisClientTest.java b/java/client/src/test/java/glide/api/RedisClientTest.java
@@ -65,6 +65,7 @@
 import static redis_request.RedisRequestOuterClass.RequestType.ConfigRewrite;
 import static redis_request.RedisRequestOuterClass.RequestType.ConfigSet;
 import static redis_request.RedisRequestOuterClass.RequestType.CustomCommand;
+import static redis_request.RedisRequestOuterClass.RequestType.DBSize;
 import static redis_request.RedisRequestOuterClass.RequestType.Decr;
 import static redis_request.RedisRequestOuterClass.RequestType.DecrBy;
 import static redis_request.RedisRequestOuterClass.RequestType.Del;
@@ -4377,6 +4378,26 @@ public void lolwut_with_version_and_params_returns_success() {
         assertEquals(value, response.get());
     }
 
+    @SneakyThrows
+    @Test
+    public void dbsize_returns_success() {
+        // setup
+        Long value = 10L;
+        CompletableFuture<Long> testResponse = new CompletableFuture<>();
+        testResponse.complete(value);
+
+        // match on protobuf request
+        when(commandManager.<Long>submitNewCommand(eq(DBSize), eq(new String[0]), any()))
+                .thenReturn(testResponse);
+
+        // exercise
+        CompletableFuture<Long> response = service.dbsize();
+
+        // verify
+        assertEquals(testResponse, response);
+        assertEquals(value, response.get());
+    }
+
     @SneakyThrows
     @Test
     public void linsert_returns_success() {

diff --git a/java/client/src/test/java/glide/api/RedisClusterClientTest.java b/java/client/src/test/java/glide/api/RedisClusterClientTest.java
@@ -23,6 +23,7 @@
 import static redis_request.RedisRequestOuterClass.RequestType.ConfigResetStat;
 import static redis_request.RedisRequestOuterClass.RequestType.ConfigRewrite;
 import static redis_request.RedisRequestOuterClass.RequestType.ConfigSet;
+import static redis_request.RedisRequestOuterClass.RequestType.DBSize;
 import static redis_request.RedisRequestOuterClass.RequestType.Echo;
 import static redis_request.RedisRequestOuterClass.RequestType.FlushAll;
 import static redis_request.RedisRequestOuterClass.RequestType.FunctionDelete;
@@ -1084,6 +1085,47 @@ public void lolwut_with_version_and_params_and_route_returns_success() {
         assertEquals(value, response.get());
     }
 
+    @SneakyThrows
+    @Test
+    public void dbsize_returns_success() {
+        // setup
+        Long value = 10L;
+        CompletableFuture<Long> testResponse = new CompletableFuture<>();
+        testResponse.complete(value);
+
+        // match on protobuf request
+        when(commandManager.<Long>submitNewCommand(eq(DBSize), eq(new String[0]), any()))
+                .thenReturn(testResponse);
+
+        // exercise
+        CompletableFuture<Long> response = service.dbsize();
+
+        // verify
+        assertEquals(testResponse, response);
+        assertEquals(value, response.get());
+    }
+
+    @SneakyThrows
+    @Test
+    public void dbsize_with_route_returns_success() {
+        // setup
+        Long value = 10L;
+        CompletableFuture<Long> testResponse = new CompletableFuture<>();
+        testResponse.complete(value);
+
+        // match on protobuf request
+        when(commandManager.<Long>submitNewCommand(
+                        eq(DBSize), eq(new String[0]), eq(ALL_PRIMARIES), any()))
+                .thenReturn(testResponse);
+
+        // exercise
+        CompletableFuture<Long> response = service.dbsize(ALL_PRIMARIES);
+
+        // verify
+        assertEquals(testResponse, response);
+        assertEquals(value, response.get());
+    }
+
     @SneakyThrows
     @Test
     public void functionLoad_returns_success() {

diff --git a/java/client/src/test/java/glide/api/models/TransactionTests.java b/java/client/src/test/java/glide/api/models/TransactionTests.java
@@ -48,6 +48,7 @@
 import static redis_request.RedisRequestOuterClass.RequestType.ConfigResetStat;
 import static redis_request.RedisRequestOuterClass.RequestType.ConfigRewrite;
 import static redis_request.RedisRequestOuterClass.RequestType.ConfigSet;
+import static redis_request.RedisRequestOuterClass.RequestType.DBSize;
 import static redis_request.RedisRequestOuterClass.RequestType.Decr;
 import static redis_request.RedisRequestOuterClass.RequestType.DecrBy;
 import static redis_request.RedisRequestOuterClass.RequestType.Del;
@@ -734,6 +735,9 @@ InfScoreBound.NEGATIVE_INFINITY, new ScoreBoundary(3, false), new Limit(1, 2)),
         results.add(Pair.of(Lolwut, buildArgs("1", "2")));
         results.add(Pair.of(Lolwut, buildArgs(VERSION_REDIS_API, "6", "42")));
 
+        transaction.dbsize();
+        results.add(Pair.of(DBSize, buildArgs()));
+
         transaction.persist("key");
         results.add(Pair.of(Persist, buildArgs("key")));
 

diff --git a/java/integTest/src/test/java/glide/TransactionTestUtilities.java b/java/integTest/src/test/java/glide/TransactionTestUtilities.java
@@ -572,7 +572,8 @@ private static Object[] serverManagementCommands(BaseTransaction<?> transaction)
                 .configResetStat()
                 .lolwut(1)
                 .flushall()
-                .flushall(ASYNC);
+                .flushall(ASYNC)
+                .dbsize();
 
         return new Object[] {
             OK, // configSet(Map.of("timeout", "1000"))
@@ -581,6 +582,7 @@ private static Object[] serverManagementCommands(BaseTransaction<?> transaction)
             "Redis ver. " + REDIS_VERSION + '\n', // lolwut(1)
             OK, // flushall()
             OK, // flushall(ASYNC)
+            0L, // dbsize()
         };
     }
 

diff --git a/java/integTest/src/test/java/glide/cluster/CommandTests.java b/java/integTest/src/test/java/glide/cluster/CommandTests.java
@@ -638,6 +638,28 @@ public void lolwut_lolwut() {
         assertTrue(clusterResponse.getSingleValue().contains("Redis ver. " + REDIS_VERSION));
     }
 
+    @Test
+    @SneakyThrows
+    public void dbsize() {
+        assertEquals(OK, clusterClient.flushall().get());
+        // dbsize should be 0 after flushall() because all keys have been deleted
+        assertEquals(0L, clusterClient.dbsize().get());
+
+        int numKeys = 10;
+        for (int i = 0; i < numKeys; i++) {
+            assertEquals(OK, clusterClient.set(UUID.randomUUID().toString(), "foo").get());
+        }
+        assertEquals(10L, clusterClient.dbsize(ALL_PRIMARIES).get());
+
+        // test dbsize with routing - flush the database first to ensure the set() call is directed to a
+        // node with 0 keys.
+        assertEquals(OK, clusterClient.flushall().get());
+        assertEquals(0L, clusterClient.dbsize().get());
+        String key = UUID.randomUUID().toString();
+        assertEquals(OK, clusterClient.set(key, "foo").get());
+        assertEquals(1L, clusterClient.dbsize(new SlotKeyRoute(key, PRIMARY)).get());
+    }
+
     @Test
     @SneakyThrows
     public void objectFreq() {

diff --git a/java/integTest/src/test/java/glide/standalone/CommandTests.java b/java/integTest/src/test/java/glide/standalone/CommandTests.java
@@ -168,9 +168,9 @@ public void move() {
         String value1 = UUID.randomUUID().toString();
         String value2 = UUID.randomUUID().toString();
         String nonExistingKey = UUID.randomUUID().toString();
-        assertEquals(false, regularClient.move(nonExistingKey, 1L).get());
-
         assertEquals(OK, regularClient.select(0).get());
+
+        assertEquals(false, regularClient.move(nonExistingKey, 1L).get());
         assertEquals(OK, regularClient.set(key1, value1).get());
         assertEquals(OK, regularClient.set(key2, value2).get());
         assertEquals(true, regularClient.move(key1, 1L).get());
@@ -344,6 +344,22 @@ public void lolwut_lolwut() {
         assertTrue(response.contains("Redis ver. " + REDIS_VERSION));
     }
 
+    @Test
+    @SneakyThrows
+    public void dbsize() {
+        assertEquals(OK, regularClient.flushall().get());
+        assertEquals(OK, regularClient.select(0).get());
+
+        int numKeys = 10;
+        for (int i = 0; i < numKeys; i++) {
+            assertEquals(OK, regularClient.set(UUID.randomUUID().toString(), "foo").get());
+        }
+        assertEquals(10L, regularClient.dbsize().get());
+
+        assertEquals(OK, regularClient.select(1).get());
+        assertEquals(0L, regularClient.dbsize().get());
+    }
+
     @Test
     @SneakyThrows
     public void objectFreq() {