Skip to content

Commit ff0cd5a

Browse files
committed
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
1 parent bdcdc97 commit ff0cd5a

File tree

2 files changed

+84
-7
lines changed

2 files changed

+84
-7
lines changed

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

+62-6
Original file line numberDiff line numberDiff line change
@@ -37,11 +37,13 @@
3737
import io.lettuce.core.protocol.Command;
3838
import io.lettuce.core.protocol.CommandArgs;
3939
import io.lettuce.core.protocol.CommandType;
40+
import io.lettuce.core.protocol.ProtocolKeyword;
4041
import io.lettuce.core.pubsub.StatefulRedisPubSubConnection;
4142
import io.lettuce.core.sentinel.api.StatefulRedisSentinelConnection;
4243
import lombok.RequiredArgsConstructor;
4344

4445
import java.lang.reflect.Constructor;
46+
import java.nio.charset.StandardCharsets;
4547
import java.util.ArrayList;
4648
import java.util.Collections;
4749
import java.util.HashMap;
@@ -405,7 +407,7 @@ public Object execute(String command, @Nullable CommandOutput commandOutputTypeH
405407
try {
406408

407409
String name = command.trim().toUpperCase();
408-
CommandType commandType = CommandType.valueOf(name);
410+
ProtocolKeyword commandType = getCommandType(name);
409411

410412
validateCommandIfRunningInTransactionMode(commandType, args);
411413

@@ -1046,14 +1048,14 @@ io.lettuce.core.ScanCursor getScanCursor(long cursorId) {
10461048
return io.lettuce.core.ScanCursor.of(Long.toString(cursorId));
10471049
}
10481050

1049-
private void validateCommandIfRunningInTransactionMode(CommandType cmd, byte[]... args) {
1051+
private void validateCommandIfRunningInTransactionMode(ProtocolKeyword cmd, byte[]... args) {
10501052

10511053
if (this.isQueueing()) {
10521054
validateCommand(cmd, args);
10531055
}
10541056
}
10551057

1056-
private void validateCommand(CommandType cmd, @Nullable byte[]... args) {
1058+
private void validateCommand(ProtocolKeyword cmd, @Nullable byte[]... args) {
10571059

10581060
RedisCommand redisCommand = RedisCommand.failsafeCommandLookup(cmd.name());
10591061
if (!RedisCommand.UNKNOWN.equals(redisCommand) && redisCommand.requiresArguments()) {
@@ -1106,6 +1108,15 @@ LettuceConnectionProvider getConnectionProvider() {
11061108
return connectionProvider;
11071109
}
11081110

1111+
private static ProtocolKeyword getCommandType(String name) {
1112+
1113+
try {
1114+
return CommandType.valueOf(name);
1115+
} catch (IllegalArgumentException e) {
1116+
return new CustomCommandType(name);
1117+
}
1118+
}
1119+
11091120
/**
11101121
* {@link TypeHints} provide {@link CommandOutput} information for a given {@link CommandType}.
11111122
*
@@ -1114,7 +1125,7 @@ LettuceConnectionProvider getConnectionProvider() {
11141125
static class TypeHints {
11151126

11161127
@SuppressWarnings("rawtypes") //
1117-
private static final Map<CommandType, Class<? extends CommandOutput>> COMMAND_OUTPUT_TYPE_MAPPING = new HashMap<>();
1128+
private static final Map<ProtocolKeyword, Class<? extends CommandOutput>> COMMAND_OUTPUT_TYPE_MAPPING = new HashMap<>();
11181129

11191130
@SuppressWarnings("rawtypes") //
11201131
private static final Map<Class<?>, Constructor<CommandOutput>> CONSTRUCTORS = new ConcurrentHashMap<>();
@@ -1275,7 +1286,7 @@ static class TypeHints {
12751286
* @return {@link ByteArrayOutput} as default when no matching {@link CommandOutput} available.
12761287
*/
12771288
@SuppressWarnings("rawtypes")
1278-
public CommandOutput getTypeHint(CommandType type) {
1289+
public CommandOutput getTypeHint(ProtocolKeyword type) {
12791290
return getTypeHint(type, new ByteArrayOutput<>(CODEC));
12801291
}
12811292

@@ -1286,7 +1297,7 @@ public CommandOutput getTypeHint(CommandType type) {
12861297
* @return
12871298
*/
12881299
@SuppressWarnings("rawtypes")
1289-
public CommandOutput getTypeHint(CommandType type, CommandOutput defaultType) {
1300+
public CommandOutput getTypeHint(ProtocolKeyword type, CommandOutput defaultType) {
12901301

12911302
if (type == null || !COMMAND_OUTPUT_TYPE_MAPPING.containsKey(type)) {
12921303
return defaultType;
@@ -1523,4 +1534,49 @@ public void onClose(StatefulConnection<?, ?> connection) {
15231534
connection.setAutoFlushCommands(true);
15241535
}
15251536
}
1537+
1538+
/**
1539+
* @since 2.3.8
1540+
*/
1541+
static class CustomCommandType implements ProtocolKeyword {
1542+
1543+
private final String name;
1544+
1545+
CustomCommandType(String name) {
1546+
this.name = name;
1547+
}
1548+
1549+
@Override
1550+
public byte[] getBytes() {
1551+
return name.getBytes(StandardCharsets.US_ASCII);
1552+
}
1553+
1554+
@Override
1555+
public String name() {
1556+
return name;
1557+
}
1558+
1559+
@Override
1560+
public boolean equals(Object o) {
1561+
1562+
if (this == o) {
1563+
return true;
1564+
}
1565+
if (!(o instanceof CustomCommandType)) {
1566+
return false;
1567+
}
1568+
CustomCommandType that = (CustomCommandType) o;
1569+
return ObjectUtils.nullSafeEquals(name, that.name);
1570+
}
1571+
1572+
@Override
1573+
public int hashCode() {
1574+
return ObjectUtils.nullSafeHashCode(name);
1575+
}
1576+
1577+
@Override
1578+
public String toString() {
1579+
return name;
1580+
}
1581+
}
15261582
}

Diff for: src/test/java/org/springframework/data/redis/connection/lettuce/LettuceConnectionUnitTestSuite.java

+22-1
Original file line numberDiff line numberDiff line change
@@ -19,11 +19,17 @@
1919
import static org.mockito.Mockito.*;
2020

2121
import io.lettuce.core.RedisClient;
22+
import io.lettuce.core.RedisFuture;
2223
import io.lettuce.core.XAddArgs;
2324
import io.lettuce.core.api.StatefulRedisConnection;
2425
import io.lettuce.core.api.async.RedisAsyncCommands;
2526
import io.lettuce.core.api.sync.RedisCommands;
27+
import io.lettuce.core.codec.ByteArrayCodec;
2628
import io.lettuce.core.codec.RedisCodec;
29+
import io.lettuce.core.output.StatusOutput;
30+
import io.lettuce.core.protocol.AsyncCommand;
31+
import io.lettuce.core.protocol.Command;
32+
import io.lettuce.core.protocol.CommandArgs;
2733

2834
import java.lang.reflect.InvocationTargetException;
2935
import java.util.Collections;
@@ -32,8 +38,8 @@
3238
import org.junit.Test;
3339
import org.junit.runner.RunWith;
3440
import org.junit.runners.Suite;
35-
3641
import org.mockito.ArgumentCaptor;
42+
3743
import org.springframework.dao.InvalidDataAccessResourceUsageException;
3844
import org.springframework.data.redis.connection.AbstractConnectionUnitTestBase;
3945
import org.springframework.data.redis.connection.RedisServerCommands.ShutdownOption;
@@ -182,6 +188,21 @@ public void xaddShouldHonorMaxlen() {
182188

183189
assertThat(args.getValue()).extracting("maxlen").isEqualTo(100L);
184190
}
191+
192+
@Test // GH-1979
193+
public void executeShouldPassThruCustomCommands() {
194+
195+
Command<byte[], byte[], String> command = new Command<>(new LettuceConnection.CustomCommandType("FOO.BAR"),
196+
new StatusOutput<>(ByteArrayCodec.INSTANCE));
197+
AsyncCommand<byte[], byte[], String> future = new AsyncCommand<>(command);
198+
future.complete();
199+
200+
when(asyncCommandsMock.dispatch(any(), any(), any())).thenReturn((RedisFuture) future);
201+
202+
connection.execute("foo.bar", command.getOutput());
203+
204+
verify(asyncCommandsMock).dispatch(eq(command.getType()), eq(command.getOutput()), any(CommandArgs.class));
205+
}
185206
}
186207

187208
public static class LettucePipelineConnectionUnitTests extends LettuceConnectionUnitTests {

0 commit comments

Comments
 (0)