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

Author: mp911de

src/main/java/org/springframework/data/redis/connection/lettuce/LettuceConnection.java

@@ -37,10 +37,12 @@
 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 Object execute(String command, @Nullable CommandOutput commandOutputTypeH
 		try {
 
 			String name = command.trim().toUpperCase();
-			CommandType commandType = CommandType.valueOf(name);
+			ProtocolKeyword commandType = getCommandType(name);
 
 			validateCommandIfRunningInTransactionMode(commandType, args);
 
@@ -1045,14 +1047,14 @@ io.lettuce.core.ScanCursor getScanCursor(long cursorId) {
 		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 @@ LettuceConnectionProvider getConnectionProvider() {
 		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 @@ LettuceConnectionProvider getConnectionProvider() {
 	static class TypeHints {
 
 		@SuppressWarnings("rawtypes") //
-		private static final Map<CommandType, Class<? extends CommandOutput>> COMMAND_OUTPUT_TYPE_MAPPING = new HashMap<>();
+		private static final Map<ProtocolKeyword, Class<? extends CommandOutput>> COMMAND_OUTPUT_TYPE_MAPPING = new HashMap<>();
 
 		@SuppressWarnings("rawtypes") //
 		private static final Map<Class<?>, Constructor<CommandOutput>> CONSTRUCTORS = new ConcurrentHashMap<>();
@@ -1275,7 +1286,7 @@ static class TypeHints {
 		 * @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 CommandOutput getTypeHint(CommandType type) {
 		 * @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 @@ static PipeliningFlushPolicy buffered(int bufferSize) {
 
 	/**
 	 * State object associated with flushing of the currently ongoing pipeline.
-	 * 
+	 *
 	 * @author Mark Paluch
 	 * @since 2.3
 	 */
@@ -1440,7 +1451,7 @@ public interface PipeliningFlushState {
 
 	/**
 	 * Implementation to flush on each command.
-	 * 
+	 *
 	 * @author Mark Paluch
 	 * @since 2.3
 	 */
@@ -1465,7 +1476,7 @@ public void onClose(StatefulConnection<?, ?> connection) {}
 
 	/**
 	 * Implementation to flush on closing the pipeline.
-	 * 
+	 *
 	 * @author Mark Paluch
 	 * @since 2.3
 	 */
@@ -1497,7 +1508,7 @@ public void onClose(StatefulConnection<?, ?> connection) {
 
 	/**
 	 * Pipeline state for buffered flushing.
-	 * 
+	 *
 	 * @author Mark Paluch
 	 * @since 2.3
 	 */
@@ -1529,4 +1540,49 @@ public void onClose(StatefulConnection<?, ?> connection) {
 			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;
+		}
+	}
 }

src/test/java/org/springframework/data/redis/connection/lettuce/LettuceConnectionUnitTests.java

@@ -19,12 +19,18 @@
 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 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 @@ void xClaimShouldNotAddJustIdFlagToArgs() {
 			}
 
 			assertThat(ReflectionTestUtils.getField(args.getValue(), "justid")).isEqualTo(false);
-
 		}
 
 		@Test // DATAREDIS-1226
@@ -216,6 +222,21 @@ void xClaimJustIdShouldAddJustIdFlagToArgs() {
 
 			assertThat(ReflectionTestUtils.getField(args.getValue(), "justid")).isEqualTo(true);
 		}
+
+		@Test // GH-1979
+		void executeShouldPassThruCustomCommands() {
+
+			Command<byte[], byte[], String> command = new Command<>(new LettuceConnection.CustomCommandType("FOO.BAR"),
+					new StatusOutput<>(ByteArrayCodec.INSTANCE));
+			AsyncCommand<byte[], byte[], String> 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 {