From aaf3cabeb63a7a91c077d9d74a45292502fa0577 Mon Sep 17 00:00:00 2001 From: Mark Paluch Date: Tue, 23 Feb 2021 09:30:00 +0100 Subject: [PATCH] Pass-thru custom Redis commands using Lettuce. We now accept unknown custom Redis commands when using the Lettuce driver. Previously, custom commands were required to exist in Lettuce's CommandType enumeration and unknown commands (such as modules) failed to run. Closes #1979 --- .../connection/lettuce/LettuceConnection.java | 76 ++++++++++++++++--- .../lettuce/LettuceConnectionUnitTests.java | 23 +++++- 2 files changed, 88 insertions(+), 11 deletions(-) diff --git a/src/main/java/org/springframework/data/redis/connection/lettuce/LettuceConnection.java b/src/main/java/org/springframework/data/redis/connection/lettuce/LettuceConnection.java index 7fe86121a..06490aa79 100644 --- a/src/main/java/org/springframework/data/redis/connection/lettuce/LettuceConnection.java +++ b/src/main/java/org/springframework/data/redis/connection/lettuce/LettuceConnection.java @@ -37,10 +37,12 @@ import io.lettuce.core.output.*; import io.lettuce.core.protocol.Command; import io.lettuce.core.protocol.CommandArgs; import io.lettuce.core.protocol.CommandType; +import io.lettuce.core.protocol.ProtocolKeyword; import io.lettuce.core.pubsub.StatefulRedisPubSubConnection; import io.lettuce.core.sentinel.api.StatefulRedisSentinelConnection; import java.lang.reflect.Constructor; +import java.nio.charset.StandardCharsets; import java.util.ArrayList; import java.util.Collections; import java.util.HashMap; @@ -404,7 +406,7 @@ public class LettuceConnection extends AbstractRedisConnection { try { String name = command.trim().toUpperCase(); - CommandType commandType = CommandType.valueOf(name); + ProtocolKeyword commandType = getCommandType(name); validateCommandIfRunningInTransactionMode(commandType, args); @@ -1045,14 +1047,14 @@ public class LettuceConnection extends AbstractRedisConnection { return io.lettuce.core.ScanCursor.of(Long.toString(cursorId)); } - private void validateCommandIfRunningInTransactionMode(CommandType cmd, byte[]... args) { + private void validateCommandIfRunningInTransactionMode(ProtocolKeyword cmd, byte[]... args) { if (this.isQueueing()) { validateCommand(cmd, args); } } - private void validateCommand(CommandType cmd, @Nullable byte[]... args) { + private void validateCommand(ProtocolKeyword cmd, @Nullable byte[]... args) { RedisCommand redisCommand = RedisCommand.failsafeCommandLookup(cmd.name()); if (!RedisCommand.UNKNOWN.equals(redisCommand) && redisCommand.requiresArguments()) { @@ -1105,6 +1107,15 @@ public class LettuceConnection extends AbstractRedisConnection { return connectionProvider; } + private static ProtocolKeyword getCommandType(String name) { + + try { + return CommandType.valueOf(name); + } catch (IllegalArgumentException e) { + return new CustomCommandType(name); + } + } + /** * {@link TypeHints} provide {@link CommandOutput} information for a given {@link CommandType}. * @@ -1113,7 +1124,7 @@ public class LettuceConnection extends AbstractRedisConnection { static class TypeHints { @SuppressWarnings("rawtypes") // - private static final Map> COMMAND_OUTPUT_TYPE_MAPPING = new HashMap<>(); + private static final Map> COMMAND_OUTPUT_TYPE_MAPPING = new HashMap<>(); @SuppressWarnings("rawtypes") // private static final Map, Constructor> CONSTRUCTORS = new ConcurrentHashMap<>(); @@ -1275,7 +1286,7 @@ public class LettuceConnection extends AbstractRedisConnection { * @return {@link ByteArrayOutput} as default when no matching {@link CommandOutput} available. */ @SuppressWarnings("rawtypes") - public CommandOutput getTypeHint(CommandType type) { + public CommandOutput getTypeHint(ProtocolKeyword type) { return getTypeHint(type, new ByteArrayOutput<>(CODEC)); } @@ -1286,7 +1297,7 @@ public class LettuceConnection extends AbstractRedisConnection { * @return */ @SuppressWarnings("rawtypes") - public CommandOutput getTypeHint(CommandType type, CommandOutput defaultType) { + public CommandOutput getTypeHint(ProtocolKeyword type, CommandOutput defaultType) { if (type == null || !COMMAND_OUTPUT_TYPE_MAPPING.containsKey(type)) { return defaultType; @@ -1407,7 +1418,7 @@ public class LettuceConnection extends AbstractRedisConnection { /** * State object associated with flushing of the currently ongoing pipeline. - * + * * @author Mark Paluch * @since 2.3 */ @@ -1440,7 +1451,7 @@ public class LettuceConnection extends AbstractRedisConnection { /** * Implementation to flush on each command. - * + * * @author Mark Paluch * @since 2.3 */ @@ -1465,7 +1476,7 @@ public class LettuceConnection extends AbstractRedisConnection { /** * Implementation to flush on closing the pipeline. - * + * * @author Mark Paluch * @since 2.3 */ @@ -1497,7 +1508,7 @@ public class LettuceConnection extends AbstractRedisConnection { /** * Pipeline state for buffered flushing. - * + * * @author Mark Paluch * @since 2.3 */ @@ -1529,4 +1540,49 @@ public class LettuceConnection extends AbstractRedisConnection { connection.setAutoFlushCommands(true); } } + + /** + * @since 2.3.8 + */ + static class CustomCommandType implements ProtocolKeyword { + + private final String name; + + CustomCommandType(String name) { + this.name = name; + } + + @Override + public byte[] getBytes() { + return name.getBytes(StandardCharsets.US_ASCII); + } + + @Override + public String name() { + return name; + } + + @Override + public boolean equals(Object o) { + + if (this == o) { + return true; + } + if (!(o instanceof CustomCommandType)) { + return false; + } + CustomCommandType that = (CustomCommandType) o; + return ObjectUtils.nullSafeEquals(name, that.name); + } + + @Override + public int hashCode() { + return ObjectUtils.nullSafeHashCode(name); + } + + @Override + public String toString() { + return name; + } + } } diff --git a/src/test/java/org/springframework/data/redis/connection/lettuce/LettuceConnectionUnitTests.java b/src/test/java/org/springframework/data/redis/connection/lettuce/LettuceConnectionUnitTests.java index e161afd87..e7882df5b 100644 --- a/src/test/java/org/springframework/data/redis/connection/lettuce/LettuceConnectionUnitTests.java +++ b/src/test/java/org/springframework/data/redis/connection/lettuce/LettuceConnectionUnitTests.java @@ -19,12 +19,18 @@ import static org.assertj.core.api.Assertions.*; import static org.mockito.Mockito.*; import io.lettuce.core.RedisClient; +import io.lettuce.core.RedisFuture; import io.lettuce.core.XAddArgs; import io.lettuce.core.XClaimArgs; import io.lettuce.core.api.StatefulRedisConnection; import io.lettuce.core.api.async.RedisAsyncCommands; import io.lettuce.core.api.sync.RedisCommands; +import io.lettuce.core.codec.ByteArrayCodec; import io.lettuce.core.codec.RedisCodec; +import io.lettuce.core.output.StatusOutput; +import io.lettuce.core.protocol.AsyncCommand; +import io.lettuce.core.protocol.Command; +import io.lettuce.core.protocol.CommandArgs; import java.lang.reflect.InvocationTargetException; import java.time.Duration; @@ -33,6 +39,7 @@ import java.util.Collections; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.mockito.ArgumentCaptor; + import org.springframework.dao.InvalidDataAccessResourceUsageException; import org.springframework.data.redis.connection.AbstractConnectionUnitTestBase; import org.springframework.data.redis.connection.RedisServerCommands.ShutdownOption; @@ -198,7 +205,6 @@ public class LettuceConnectionUnitTests { } assertThat(ReflectionTestUtils.getField(args.getValue(), "justid")).isEqualTo(false); - } @Test // DATAREDIS-1226 @@ -216,6 +222,21 @@ public class LettuceConnectionUnitTests { assertThat(ReflectionTestUtils.getField(args.getValue(), "justid")).isEqualTo(true); } + + @Test // GH-1979 + void executeShouldPassThruCustomCommands() { + + Command command = new Command<>(new LettuceConnection.CustomCommandType("FOO.BAR"), + new StatusOutput<>(ByteArrayCodec.INSTANCE)); + AsyncCommand future = new AsyncCommand<>(command); + future.complete(); + + when(asyncCommandsMock.dispatch(any(), any(), any())).thenReturn((RedisFuture) future); + + connection.execute("foo.bar", command.getOutput()); + + verify(asyncCommandsMock).dispatch(eq(command.getType()), eq(command.getOutput()), any(CommandArgs.class)); + } } public static class LettucePipelineConnectionUnitTests extends BasicUnitTests {