Java: add invokeScript to run lua scripts (valkey-io#1160)

* Java: add invokeScript to run lua scripts (#151) * Initial invokeScript() command Signed-off-by: Andrew Carbonetto <[email protected]> * Add lua script eval to Java Signed-off-by: Andrew Carbonetto <[email protected]> --------- Signed-off-by: Andrew Carbonetto <[email protected]> * Update javadoc comments Signed-off-by: Andrew Carbonetto <[email protected]> * Add finalize() to Script Signed-off-by: Andrew Carbonetto <[email protected]> * Make finalize protected Signed-off-by: Andrew Carbonetto <[email protected]> * Rename function; clean up javadoc Signed-off-by: Andrew Carbonetto <[email protected]> * Update examples in invokeScript Signed-off-by: Andrew Carbonetto <[email protected]> * Clean up test Signed-off-by: Andrew Carbonetto <[email protected]> --------- Signed-off-by: Andrew Carbonetto <[email protected]>
Bit-Quill · Apr 1, 2024 · 8207e3a · 8207e3a
1 parent 4151b15
commit 8207e3a
Show file tree

Hide file tree

Showing 11 changed files with 324 additions and 1 deletion.
diff --git a/java/client/src/main/java/glide/api/BaseClient.java b/java/client/src/main/java/glide/api/BaseClient.java
@@ -58,7 +58,9 @@
 import glide.api.commands.SetBaseCommands;
 import glide.api.commands.SortedSetBaseCommands;
 import glide.api.commands.StringCommands;
+import glide.api.models.Script;
 import glide.api.models.commands.ExpireOptions;
+import glide.api.models.commands.ScriptOptions;
 import glide.api.models.commands.SetOptions;
 import glide.api.models.commands.ZaddOptions;
 import glide.api.models.configuration.BaseClientConfiguration;
@@ -72,6 +74,7 @@
 import glide.managers.BaseCommandResponseResolver;
 import glide.managers.CommandManager;
 import glide.managers.ConnectionManager;
+import java.util.List;
 import java.util.Map;
 import java.util.Set;
 import java.util.concurrent.CompletableFuture;
@@ -529,6 +532,19 @@ public CompletableFuture<Long> ttl(@NonNull String key) {
         return commandManager.submitNewCommand(TTL, new String[] {key}, this::handleLongResponse);
     }
 
+    @Override
+    public CompletableFuture<Object> invokeScript(@NonNull Script script) {
+        return commandManager.submitScript(
+                script, List.of(), List.of(), this::handleObjectOrNullResponse);
+    }
+
+    @Override
+    public CompletableFuture<Object> invokeScript(
+            @NonNull Script script, @NonNull ScriptOptions options) {
+        return commandManager.submitScript(
+                script, options.getKeys(), options.getArgs(), this::handleObjectOrNullResponse);
+    }
+
     @Override
     public CompletableFuture<Long> zadd(
             @NonNull String key,

diff --git a/java/client/src/main/java/glide/api/commands/GenericBaseCommands.java b/java/client/src/main/java/glide/api/commands/GenericBaseCommands.java
@@ -1,7 +1,9 @@
 /** Copyright GLIDE-for-Redis Project Contributors - SPDX Identifier: Apache-2.0 */
 package glide.api.commands;
 
+import glide.api.models.Script;
 import glide.api.models.commands.ExpireOptions;
+import glide.api.models.commands.ScriptOptions;
 import java.util.concurrent.CompletableFuture;
 
 /**
@@ -270,6 +272,53 @@ CompletableFuture<Boolean> pexpireAt(
      */
     CompletableFuture<Long> ttl(String key);
 
+    /**
+     * Invokes a Lua script.<br>
+     * This method simplifies the process of invoking scripts on a Redis server by using an object
+     * that represents a Lua script. The script loading and execution will all be handled internally.
+     * If the script has not already been loaded, it will be loaded automatically using the Redis
+     * <code>SCRIPT LOAD</code> command. After that, it will be invoked using the Redis <code>EVALSHA
+     * </code> command.
+     *
+     * @see <a href="https://redis.io/commands/script-load/">SCRIPT LOAD</a> and <a
+     *     href="https://redis.io/commands/evalsha/">EVALSHA</a> for details.
+     * @param script The Lua script to execute.
+     * @return a value that depends on the script that was executed.
+     * @example
+     *     <pre>{@code
+     * try(Script luaScript = new Script("return 'Hello'")) {
+     *     String result = (String) client.invokeScript(luaScript).get();
+     *     assert result.equals("Hello");
+     * }
+     * }</pre>
+     */
+    CompletableFuture<Object> invokeScript(Script script);
+
+    /**
+     * Invokes a Lua script with its keys and arguments.<br>
+     * This method simplifies the process of invoking scripts on a Redis server by using an object
+     * that represents a Lua script. The script loading, argument preparation, and execution will all
+     * be handled internally. If the script has not already been loaded, it will be loaded
+     * automatically using the Redis <code>SCRIPT LOAD</code> command. After that, it will be invoked
+     * using the Redis <code>EVALSHA</code> command.
+     *
+     * @see <a href="https://redis.io/commands/script-load/">SCRIPT LOAD</a> and <a
+     *     href="https://redis.io/commands/evalsha/">EVALSHA</a> for details.
+     * @param script The Lua script to execute.
+     * @param options The script option that contains keys and arguments for the script.
+     * @return a value that depends on the script that was executed.
+     * @example
+     *     <pre>{@code
+     * try(Script luaScript = new Script("return { KEYS[1], ARGV[1] }")) {
+     *     ScriptOptions scriptOptions = ScriptOptions.builder().key("foo").arg("bar").build();
+     *     Object[] result = (Object[]) client.invokeScript(luaScript, scriptOptions).get();
+     *     assert result[0].equals("foo");
+     *     assert result[1].equals("bar");
+     * }
+     * }</pre>
+     */
+    CompletableFuture<Object> invokeScript(Script script, ScriptOptions options);
+
     /**
      * Returns the remaining time to live of <code>key</code> that has a timeout, in milliseconds.
      *

diff --git a/java/client/src/main/java/glide/api/models/Script.java b/java/client/src/main/java/glide/api/models/Script.java
@@ -0,0 +1,45 @@
+/** Copyright GLIDE-for-Redis Project Contributors - SPDX Identifier: Apache-2.0 */
+package glide.api.models;
+
+import static glide.ffi.resolvers.ScriptResolver.dropScript;
+import static glide.ffi.resolvers.ScriptResolver.storeScript;
+
+import glide.api.commands.GenericBaseCommands;
+import lombok.Getter;
+
+/**
+ * A wrapper for a Script object for {@link GenericBaseCommands#invokeScript(Script)} As long as
+ * this object is not closed, the script's code is saved in memory, and can be resent to the server.
+ * Script should be enclosed with a try-with-resource block or {@link Script#close()} must be called
+ * to invalidate the code hash.
+ */
+public class Script implements AutoCloseable {
+
+    /** Hash string representing the code. */
+    @Getter private final String hash;
+
+    /**
+     * Wraps around creating a Script object from <code>code</code>.
+     *
+     * @param code To execute with a ScriptInvoke call.
+     */
+    public Script(String code) {
+        hash = storeScript(code);
+    }
+
+    /** Drop the linked script from glide-rs <code>code</code>. */
+    @Override
+    public void close() throws Exception {
+        dropScript(hash);
+    }
+
+    @Override
+    protected void finalize() throws Throwable {
+        try {
+            // Drop the linked script on garbage collection.
+            this.close();
+        } finally {
+            super.finalize();
+        }
+    }
+}
diff --git a/java/client/src/main/java/glide/api/models/commands/ScriptOptions.java b/java/client/src/main/java/glide/api/models/commands/ScriptOptions.java
@@ -0,0 +1,24 @@
+/** Copyright GLIDE-for-Redis Project Contributors - SPDX Identifier: Apache-2.0 */
+package glide.api.models.commands;
+
+import glide.api.commands.GenericBaseCommands;
+import glide.api.models.Script;
+import java.util.List;
+import lombok.Builder;
+import lombok.Getter;
+import lombok.Singular;
+
+/**
+ * Optional arguments for {@link GenericBaseCommands#invokeScript(Script, ScriptOptions)} command.
+ *
+ * @see <a href="https://redis.io/commands/evalsha/">redis.io</a>
+ */
+@Builder
+public final class ScriptOptions {
+
+    /** The keys that are used in the script. */
+    @Singular @Getter private final List<String> keys;
+
+    /** The arguments for the script. */
+    @Singular @Getter private final List<String> args;
+}
diff --git a/java/client/src/main/java/glide/ffi/resolvers/ScriptResolver.java b/java/client/src/main/java/glide/ffi/resolvers/ScriptResolver.java
@@ -0,0 +1,25 @@
+/** Copyright GLIDE-for-Redis Project Contributors - SPDX Identifier: Apache-2.0 */
+package glide.ffi.resolvers;
+
+public class ScriptResolver {
+
+    // TODO: consider lazy loading the glide_rs library
+    static {
+        System.loadLibrary("glide_rs");
+    }
+
+    /**
+     * Loads a Lua script into the scripts cache, without executing it.
+     *
+     * @param code The Lua script
+     * @return String representing the saved hash
+     */
+    public static native String storeScript(String code);
+
+    /**
+     * Unload or drop the stored Lua script from the script cache.
+     *
+     * @param hash
+     */
+    public static native void dropScript(String hash);
+}
diff --git a/java/client/src/main/java/glide/managers/CommandManager.java b/java/client/src/main/java/glide/managers/CommandManager.java
@@ -2,6 +2,7 @@
 package glide.managers;
 
 import glide.api.models.ClusterTransaction;
+import glide.api.models.Script;
 import glide.api.models.Transaction;
 import glide.api.models.configuration.RequestRoutingConfiguration.ByAddressRoute;
 import glide.api.models.configuration.RequestRoutingConfiguration.Route;
@@ -12,6 +13,7 @@
 import glide.api.models.exceptions.RequestException;
 import glide.connectors.handlers.CallbackDispatcher;
 import glide.connectors.handlers.ChannelHandler;
+import java.util.List;
 import java.util.Optional;
 import java.util.concurrent.CompletableFuture;
 import lombok.RequiredArgsConstructor;
@@ -21,6 +23,7 @@
 import redis_request.RedisRequestOuterClass.RedisRequest;
 import redis_request.RedisRequestOuterClass.RequestType;
 import redis_request.RedisRequestOuterClass.Routes;
+import redis_request.RedisRequestOuterClass.ScriptInvocation;
 import redis_request.RedisRequestOuterClass.SimpleRoutes;
 import redis_request.RedisRequestOuterClass.SlotTypes;
 import response.ResponseOuterClass.Response;
@@ -85,6 +88,25 @@ public <T> CompletableFuture<T> submitNewCommand(
         return submitCommandToChannel(command, responseHandler);
     }
 
+    /**
+     * Build a Script (by hash) request to send to Redis.
+     *
+     * @param script Lua script hash object
+     * @param keys The keys that are used in the script
+     * @param args The arguments for the script
+     * @param responseHandler The handler for the response object
+     * @return A result promise of type T
+     */
+    public <T> CompletableFuture<T> submitScript(
+            Script script,
+            List<String> keys,
+            List<String> args,
+            RedisExceptionCheckedFunction<Response, T> responseHandler) {
+
+        RedisRequest.Builder command = prepareRedisRequest(script, keys, args);
+        return submitCommandToChannel(command, responseHandler);
+    }
+
     /**
      * Build a Transaction and send.
      *
@@ -168,6 +190,29 @@ protected RedisRequest.Builder prepareRedisRequest(Transaction transaction) {
         return builder;
     }
 
+    /**
+     * Build a protobuf Script Invoke request.
+     *
+     * @param script Redis Script
+     * @param keys keys for the Script
+     * @param args args for the Script
+     * @return An uncompleted request. {@link CallbackDispatcher} is responsible to complete it by
+     *     adding a callback id.
+     */
+    protected RedisRequest.Builder prepareRedisRequest(
+            Script script, List<String> keys, List<String> args) {
+        RedisRequest.Builder builder =
+                RedisRequest.newBuilder()
+                        .setScriptInvocation(
+                                ScriptInvocation.newBuilder()
+                                        .setHash(script.getHash())
+                                        .addAllKeys(keys)
+                                        .addAllArgs(args)
+                                        .build());
+
+        return builder;
+    }
+
     /**
      * Build a protobuf transaction request object with routing options.
      *

diff --git a/java/client/src/test/java/glide/api/RedisClientTest.java b/java/client/src/test/java/glide/api/RedisClientTest.java
@@ -73,16 +73,20 @@
 import static redis_request.RedisRequestOuterClass.RequestType.Zcard;
 import static redis_request.RedisRequestOuterClass.RequestType.Zrem;
 
+import glide.api.models.Script;
 import glide.api.models.commands.ExpireOptions;
 import glide.api.models.commands.InfoOptions;
+import glide.api.models.commands.ScriptOptions;
 import glide.api.models.commands.SetOptions;
 import glide.api.models.commands.SetOptions.Expiry;
 import glide.api.models.commands.ZaddOptions;
 import glide.managers.CommandManager;
 import glide.managers.ConnectionManager;
 import java.util.LinkedHashMap;
+import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.UUID;
 import java.util.concurrent.CompletableFuture;
 import lombok.SneakyThrows;
 import org.apache.commons.lang3.ArrayUtils;
@@ -580,6 +584,58 @@ public void ttl_returns_success() {
         assertEquals(ttl, response.get());
     }
 
+    @SneakyThrows
+    @Test
+    public void invokeScript_returns_success() {
+        // setup
+        Script script = mock(Script.class);
+        String hash = UUID.randomUUID().toString();
+        when(script.getHash()).thenReturn(hash);
+        String payload = "hello";
+
+        CompletableFuture<Object> testResponse = new CompletableFuture<>();
+        testResponse.complete(payload);
+
+        // match on protobuf request
+        when(commandManager.<Object>submitScript(eq(script), eq(List.of()), eq(List.of()), any()))
+                .thenReturn(testResponse);
+
+        // exercise
+        CompletableFuture<Object> response = service.invokeScript(script);
+
+        // verify
+        assertEquals(testResponse, response);
+        assertEquals(payload, response.get());
+    }
+
+    @SneakyThrows
+    @Test
+    public void invokeScript_with_ScriptOptions_returns_success() {
+        // setup
+        Script script = mock(Script.class);
+        String hash = UUID.randomUUID().toString();
+        when(script.getHash()).thenReturn(hash);
+        String payload = "hello";
+
+        ScriptOptions options =
+                ScriptOptions.builder().key("key1").key("key2").arg("arg1").arg("arg2").build();
+
+        CompletableFuture<Object> testResponse = new CompletableFuture<>();
+        testResponse.complete(payload);
+
+        // match on protobuf request
+        when(commandManager.<Object>submitScript(
+                        eq(script), eq(List.of("key1", "key2")), eq(List.of("arg1", "arg2")), any()))
+                .thenReturn(testResponse);
+
+        // exercise
+        CompletableFuture<Object> response = service.invokeScript(script, options);
+
+        // verify
+        assertEquals(testResponse, response);
+        assertEquals(payload, response.get());
+    }
+
     @SneakyThrows
     @Test
     public void pttl_returns_success() {