Implement OBJECT REFCOUNT command

Bit-Quill · Apr 22, 2024 · 4447db7 · 4447db7
1 parent e9f7702
commit 4447db7
Show file tree

Hide file tree

Showing 10 changed files with 114 additions and 0 deletions.
diff --git a/glide-core/src/protobuf/redis_request.proto b/glide-core/src/protobuf/redis_request.proto
@@ -163,6 +163,7 @@ enum RequestType {
     GeoAdd = 121;
     GeoHash = 122;
     ObjectEncoding = 123;
+    ObjectRefcount = 126;
 }
 
 message Command {

diff --git a/glide-core/src/request_type.rs b/glide-core/src/request_type.rs
@@ -131,6 +131,7 @@ pub enum RequestType {
     GeoAdd = 121,
     GeoHash = 122,
     ObjectEncoding = 123,
+    ObjectRefcount = 126,
 }
 
 fn get_two_word_command(first: &str, second: &str) -> Cmd {
@@ -265,6 +266,7 @@ impl From<::protobuf::EnumOrUnknown<ProtobufRequestType>> for RequestType {
             ProtobufRequestType::GeoAdd => RequestType::GeoAdd,
             ProtobufRequestType::GeoHash => RequestType::GeoHash,
             ProtobufRequestType::ObjectEncoding => RequestType::ObjectEncoding,
+            ProtobufRequestType::ObjectRefcount => RequestType::ObjectRefcount,
         }
     }
 }
@@ -395,6 +397,7 @@ impl RequestType {
             RequestType::GeoAdd => Some(cmd("GEOADD")),
             RequestType::GeoHash => Some(cmd("GEOHASH")),
             RequestType::ObjectEncoding => Some(get_two_word_command("OBJECT", "ENCODING")),
+            RequestType::ObjectRefcount => Some(get_two_word_command("OBJECT", "REFCOUNT")),
         }
     }
 }
diff --git a/java/client/src/main/java/glide/api/BaseClient.java b/java/client/src/main/java/glide/api/BaseClient.java
@@ -42,6 +42,7 @@
 import static redis_request.RedisRequestOuterClass.RequestType.MGet;
 import static redis_request.RedisRequestOuterClass.RequestType.MSet;
 import static redis_request.RedisRequestOuterClass.RequestType.ObjectEncoding;
+import static redis_request.RedisRequestOuterClass.RequestType.ObjectRefcount;
 import static redis_request.RedisRequestOuterClass.RequestType.PExpire;
 import static redis_request.RedisRequestOuterClass.RequestType.PExpireAt;
 import static redis_request.RedisRequestOuterClass.RequestType.PTTL;
@@ -341,6 +342,12 @@ public CompletableFuture<String> objectEncoding(@NonNull String key) {
                 ObjectEncoding, new String[] {key}, this::handleStringOrNullResponse);
     }
 
+    @Override
+    public CompletableFuture<Long> objectRefcount(@NonNull String key) {
+        return commandManager.submitNewCommand(
+                ObjectRefcount, new String[] {key}, this::handleLongOrNullResponse);
+    }
+
     @Override
     public CompletableFuture<Long> incr(@NonNull String key) {
         return commandManager.submitNewCommand(Incr, new String[] {key}, this::handleLongResponse);

diff --git a/java/client/src/main/java/glide/api/commands/GenericBaseCommands.java b/java/client/src/main/java/glide/api/commands/GenericBaseCommands.java
@@ -389,4 +389,22 @@ CompletableFuture<Boolean> pexpireAt(
      * }</pre>
      */
     CompletableFuture<String> objectEncoding(String key);
+
+    /**
+     * Returns the reference count of the object stored at <code>key</code>.
+     *
+     * @see <a href="https://redis.io/commands/object-refcount/">redis.io</a> for details.
+     * @param key The <code>key</code> of the object to get the reference count of.
+     * @return If <code>key</code> exists, returns the reference count of the object stored at <code>
+     *     key</code> as a <code>Long</code>. Otherwise, returns <code>null</code>.
+     * @example
+     *     <pre>{@code
+     * Long refcount = client.objectRefcount("my_hash").get();
+     * assert refcount == 2L;
+     *
+     * refcount = client.objectRefcount("non_existing_key").get();
+     * assert refcount == null;
+     * }</pre>
+     */
+    CompletableFuture<Long> objectRefcount(String key);
 }
diff --git a/java/client/src/main/java/glide/api/models/BaseTransaction.java b/java/client/src/main/java/glide/api/models/BaseTransaction.java
@@ -53,6 +53,7 @@
 import static redis_request.RedisRequestOuterClass.RequestType.MGet;
 import static redis_request.RedisRequestOuterClass.RequestType.MSet;
 import static redis_request.RedisRequestOuterClass.RequestType.ObjectEncoding;
+import static redis_request.RedisRequestOuterClass.RequestType.ObjectRefcount;
 import static redis_request.RedisRequestOuterClass.RequestType.PExpire;
 import static redis_request.RedisRequestOuterClass.RequestType.PExpireAt;
 import static redis_request.RedisRequestOuterClass.RequestType.PTTL;
@@ -2240,6 +2241,21 @@ public T objectEncoding(@NonNull String key) {
         return getThis();
     }
 
+    /**
+     * Returns the reference count of the object stored at <code>key</code>.
+     *
+     * @see <a href="https://redis.io/commands/object-refcount/">redis.io</a> for details.
+     * @param key The <code>key</code> of the object to get the reference count of.
+     * @return Command response - If <code>key</code> exists, returns the reference count of the
+     *     object stored at <code>key</code> as a <code>Long</code>. Otherwise, returns <code>null
+     *     </code>.
+     */
+    public T objectRefcount(@NonNull String key) {
+        ArgsArray commandArgs = buildArgs(key);
+        protobufTransaction.addCommands(buildCommand(ObjectRefcount, commandArgs));
+        return getThis();
+    }
+
     /** Build protobuf {@link Command} object for given command and arguments. */
     protected Command buildCommand(RequestType requestType) {
         return buildCommand(requestType, buildArgs());

diff --git a/java/client/src/test/java/glide/api/RedisClientTest.java b/java/client/src/test/java/glide/api/RedisClientTest.java
@@ -73,6 +73,7 @@
 import static redis_request.RedisRequestOuterClass.RequestType.MGet;
 import static redis_request.RedisRequestOuterClass.RequestType.MSet;
 import static redis_request.RedisRequestOuterClass.RequestType.ObjectEncoding;
+import static redis_request.RedisRequestOuterClass.RequestType.ObjectRefcount;
 import static redis_request.RedisRequestOuterClass.RequestType.PExpire;
 import static redis_request.RedisRequestOuterClass.RequestType.PExpireAt;
 import static redis_request.RedisRequestOuterClass.RequestType.PTTL;
@@ -3282,4 +3283,26 @@ public void objectEncoding_returns_success() {
         assertEquals(testResponse, response);
         assertEquals(encoding, payload);
     }
+
+    @SneakyThrows
+    @Test
+    public void objectRefcount_returns_success() {
+        // setup
+        String key = "testKey";
+        Long refcount = 0L;
+        CompletableFuture<Long> testResponse = new CompletableFuture<>();
+        testResponse.complete(refcount);
+
+        // match on protobuf request
+        when(commandManager.<Long>submitNewCommand(eq(ObjectRefcount), eq(new String[] {key}), any()))
+                .thenReturn(testResponse);
+
+        // exercise
+        CompletableFuture<Long> response = service.objectRefcount(key);
+        Long payload = response.get();
+
+        // verify
+        assertEquals(testResponse, response);
+        assertEquals(refcount, payload);
+    }
 }
diff --git a/java/client/src/test/java/glide/api/models/TransactionTests.java b/java/client/src/test/java/glide/api/models/TransactionTests.java
@@ -58,6 +58,7 @@
 import static redis_request.RedisRequestOuterClass.RequestType.MGet;
 import static redis_request.RedisRequestOuterClass.RequestType.MSet;
 import static redis_request.RedisRequestOuterClass.RequestType.ObjectEncoding;
+import static redis_request.RedisRequestOuterClass.RequestType.ObjectRefcount;
 import static redis_request.RedisRequestOuterClass.RequestType.PExpire;
 import static redis_request.RedisRequestOuterClass.RequestType.PExpireAt;
 import static redis_request.RedisRequestOuterClass.RequestType.PTTL;
@@ -512,6 +513,9 @@ InfScoreBound.NEGATIVE_INFINITY, new ScoreBoundary(3, false), new Limit(1, 2)),
         transaction.objectEncoding("key");
         results.add(Pair.of(ObjectEncoding, buildArgs("key")));
 
+        transaction.objectRefcount("key");
+        results.add(Pair.of(ObjectRefcount, buildArgs("key")));
+
         var protobufTransaction = transaction.getProtobufTransaction().build();
 
         for (int idx = 0; idx < protobufTransaction.getCommandsCount(); idx++) {

diff --git a/java/integTest/src/test/java/glide/SharedCommandTests.java b/java/integTest/src/test/java/glide/SharedCommandTests.java
@@ -2715,4 +2715,21 @@ public void objectEncoding_returns_stream(BaseClient client) {
         assertNotNull(client.xadd(streamKey, Map.of("field", "value")));
         assertEquals("stream", client.objectEncoding(streamKey).get());
     }
+
+    @SneakyThrows
+    @ParameterizedTest(autoCloseArguments = false)
+    @MethodSource("getClients")
+    public void objectRefcount_returns_null(BaseClient client) {
+        String nonExistingKey = UUID.randomUUID().toString();
+        assertNull(client.objectRefcount(nonExistingKey).get());
+    }
+
+    @SneakyThrows
+    @ParameterizedTest(autoCloseArguments = false)
+    @MethodSource("getClients")
+    public void objectRefcount(BaseClient client) {
+        String key = UUID.randomUUID().toString();
+        assertEquals(OK, client.set(key, "").get());
+        assertTrue(client.objectRefcount(key).get() >= 0L);
+    }
 }
diff --git a/java/integTest/src/test/java/glide/cluster/ClusterTransactionTests.java b/java/integTest/src/test/java/glide/cluster/ClusterTransactionTests.java
@@ -92,4 +92,17 @@ public void lastsave() {
         var response = clusterClient.exec(new ClusterTransaction().lastsave()).get();
         assertTrue(Instant.ofEpochSecond((long) response[0]).isAfter(yesterday));
     }
+
+    // TODO: Enable when https://github.com/amazon-contributing/redis-rs/pull/138 is merged.
+    // @Test
+    // @SneakyThrows
+    // public void objectRefcount() {
+    //    String objectRefcountKey = "key";
+    //    ClusterTransaction transaction = new ClusterTransaction();
+    //    transaction.set(objectRefcountKey, "");
+    //    transaction.objectRefcount(objectRefcountKey);
+    //    var response = clusterClient.exec(transaction).get();
+    //    assertEquals(OK, response[0]);
+    //    assertTrue((long) response[1] >= 0L);
+    // }
 }
diff --git a/java/integTest/src/test/java/glide/standalone/TransactionTests.java b/java/integTest/src/test/java/glide/standalone/TransactionTests.java
@@ -121,4 +121,16 @@ public void lastsave() {
         var response = client.exec(new Transaction().lastsave()).get();
         assertTrue(Instant.ofEpochSecond((long) response[0]).isAfter(yesterday));
     }
+
+    @Test
+    @SneakyThrows
+    public void objectRefcount() {
+        String objectRefcountKey = "key";
+        Transaction transaction = new Transaction();
+        transaction.set(objectRefcountKey, "");
+        transaction.objectRefcount(objectRefcountKey);
+        var response = client.exec(transaction).get();
+        assertEquals(OK, response[0]);
+        assertTrue((long) response[1] >= 0L);
+    }
 }