Proof-of-concept for automatic key prefixing

Author: R-J Lim

src/main/java/redis/clients/jedis/ClusterCommandArgumentsWithPrefixedKeys.java

@@ -0,0 +1,46 @@
+package redis.clients.jedis;
+
+import redis.clients.jedis.args.Rawable;
+import redis.clients.jedis.args.RawableFactory;
+import redis.clients.jedis.commands.ProtocolCommand;
+import redis.clients.jedis.util.SafeEncoder;
+
+public class ClusterCommandArgumentsWithPrefixedKeys extends ClusterCommandArguments {
+    private final byte[] prefix;
+    private final String prefixString;
+
+    public ClusterCommandArgumentsWithPrefixedKeys(ProtocolCommand command, String prefixString) {
+        super(command);
+        this.prefixString = prefixString;
+        prefix = SafeEncoder.encode(prefixString);
+    }
+
+    public CommandArguments key(Object key) {
+        return super.key(namespacedKey(key));
+    }
+
+    private Object namespacedKey(Object key) {
+        if (key instanceof Rawable) {
+            byte[] raw = ((Rawable) key).getRaw();
+            return RawableFactory.from(namespacedKeyBytes(raw));
+        }
+
+        if (key instanceof byte[]) {
+            return namespacedKeyBytes((byte[]) key);
+        }
+
+        if (key instanceof String) {
+            String raw = (String) key;
+            return prefixString + raw;
+        }
+
+        throw new IllegalArgumentException("\"" + key.toString() + "\" is not a valid argument.");
+    }
+
+    private byte[] namespacedKeyBytes(byte[] key) {
+        byte[] namespaced = new byte[prefix.length + key.length];
+        System.arraycopy(prefix, 0, namespaced, 0, prefix.length);
+        System.arraycopy(key, 0, namespaced, prefix.length, key.length);
+        return namespaced;
+    }
+}

src/main/java/redis/clients/jedis/ClusterCommandObjectsWithPrefixedKeys.java

@@ -0,0 +1,16 @@
+package redis.clients.jedis;
+
+import redis.clients.jedis.commands.ProtocolCommand;
+
+public class ClusterCommandObjectsWithPrefixedKeys extends ClusterCommandObjects {
+    // For the purposes of this demonstration, the prefix is assigned statically.
+    // Additional changes are required to prevent the parent class CommandObjects
+    // from calling commandArguments in its constructor, which would be a prerequisite
+    // to making this field into an instance field.
+    public static String PREFIX_STRING;
+
+    @Override
+    protected ClusterCommandArguments commandArguments(ProtocolCommand command) {
+        return new ClusterCommandArgumentsWithPrefixedKeys(command, PREFIX_STRING);
+    }
+}

src/main/java/redis/clients/jedis/JedisClusterWithPrefixedKeys.java

@@ -0,0 +1,61 @@
+package redis.clients.jedis;
+
+import redis.clients.jedis.executors.ClusterCommandExecutor;
+import redis.clients.jedis.providers.ClusterConnectionProvider;
+import redis.clients.jedis.util.JedisClusterCRC16;
+
+import java.time.Duration;
+import java.util.Collections;
+import java.util.Map;
+
+public class JedisClusterWithPrefixedKeys extends UnifiedJedis {
+
+    public JedisClusterWithPrefixedKeys(HostAndPort node, JedisClientConfig clientConfig) {
+        this(new ClusterConnectionProvider(Collections.singleton(node), clientConfig), clientConfig);
+    }
+
+    public JedisClusterWithPrefixedKeys(ClusterConnectionProvider provider, JedisClientConfig clientConfig) {
+        super(new ClusterCommandExecutor(provider, 5, Duration.ofSeconds(100)), provider, new ClusterCommandObjectsWithPrefixedKeys(), clientConfig.getRedisProtocol());
+
+    }
+
+    public Map<String, ConnectionPool> getClusterNodes() {
+        return ((ClusterConnectionProvider) provider).getNodes();
+    }
+
+    public Connection getConnectionFromSlot(int slot) {
+        return ((ClusterConnectionProvider) provider).getConnectionFromSlot(slot);
+    }
+
+    // commands
+    public long spublish(String channel, String message) {
+        return executeCommand(commandObjects.spublish(channel, message));
+    }
+
+    public long spublish(byte[] channel, byte[] message) {
+        return executeCommand(commandObjects.spublish(channel, message));
+    }
+
+    public void ssubscribe(final JedisShardedPubSub jedisPubSub, final String... channels) {
+        try (Connection connection = getConnectionFromSlot(JedisClusterCRC16.getSlot(channels[0]))) {
+            jedisPubSub.proceed(connection, channels);
+        }
+    }
+
+    public void ssubscribe(BinaryJedisShardedPubSub jedisPubSub, final byte[]... channels) {
+        try (Connection connection = getConnectionFromSlot(JedisClusterCRC16.getSlot(channels[0]))) {
+            jedisPubSub.proceed(connection, channels);
+        }
+    }
+    // commands
+
+    @Override
+    public ClusterPipeline pipelined() {
+        return new ClusterPipeline((ClusterConnectionProvider) provider, (ClusterCommandObjects) commandObjects);
+    }
+
+    @Override
+    public Transaction multi() {
+        throw new UnsupportedOperationException();
+    }
+}

src/main/java/redis/clients/jedis/UnifiedJedis.java

@@ -206,12 +206,12 @@ public UnifiedJedis(CommandExecutor executor) {
     this(executor, (ConnectionProvider) null);
   }
 
-  private UnifiedJedis(CommandExecutor executor, ConnectionProvider provider) {
+  public UnifiedJedis(CommandExecutor executor, ConnectionProvider provider) {
     this(executor, provider, new CommandObjects());
   }
 
   // Uses a fetched connection to process protocol. Should be avoided if possible.
-  private UnifiedJedis(CommandExecutor executor, ConnectionProvider provider, CommandObjects commandObjects) {
+  public UnifiedJedis(CommandExecutor executor, ConnectionProvider provider, CommandObjects commandObjects) {
     this(executor, provider, commandObjects, null);
     if (this.provider != null) {
       try (Connection conn = this.provider.getConnection()) {
@@ -223,7 +223,7 @@ private UnifiedJedis(CommandExecutor executor, ConnectionProvider provider, Comm
     }
   }
 
-  private UnifiedJedis(CommandExecutor executor, ConnectionProvider provider, CommandObjects commandObjects,
+  public UnifiedJedis(CommandExecutor executor, ConnectionProvider provider, CommandObjects commandObjects,
       RedisProtocol protocol) {
     this.provider = provider;
     this.executor = executor;

src/test/java/redis/clients/jedis/JedisClusterWithPrefixedKeysTest.java

@@ -0,0 +1,38 @@
+package redis.clients.jedis;
+
+import org.junit.Test;
+import redis.clients.jedis.resps.Tuple;
+
+import java.nio.charset.StandardCharsets;
+import java.util.Collections;
+
+import static junit.framework.TestCase.assertEquals;
+
+public class JedisClusterWithPrefixedKeysTest extends JedisClusterTestBase {
+    private static final DefaultJedisClientConfig DEFAULT_CLIENT_CONFIG = DefaultJedisClientConfig.builder().password("cluster").build();
+    private static final int DEFAULT_TIMEOUT = 2000;
+    private static final int DEFAULT_REDIRECTIONS = 5;
+    private static final ConnectionPoolConfig DEFAULT_POOL_CONFIG = new ConnectionPoolConfig();
+
+    @Test
+    public void hasPrefixedKeys() {
+        HostAndPort hp = new HostAndPort("127.0.0.1", 7379);
+        ClusterCommandObjectsWithPrefixedKeys.PREFIX_STRING = "test-prefix:";
+
+        try (JedisClusterWithPrefixedKeys cluster = new JedisClusterWithPrefixedKeys(hp, DEFAULT_CLIENT_CONFIG)) {
+            cluster.set("foo1", "bar1");
+            cluster.set("foo2".getBytes(StandardCharsets.UTF_8), "bar2".getBytes(StandardCharsets.UTF_8));
+            ClusterPipeline pipeline = cluster.pipelined();
+            pipeline.incr("foo3");
+            pipeline.zadd("foo4", 1234, "bar4");
+            pipeline.sync();
+        }
+
+        try (JedisCluster cluster = new JedisCluster(hp, DEFAULT_CLIENT_CONFIG, DEFAULT_REDIRECTIONS, DEFAULT_POOL_CONFIG)) {
+            assertEquals("bar1", cluster.get("test-prefix:foo1"));
+            assertEquals("bar2", cluster.get("test-prefix:foo2"));
+            assertEquals("1", cluster.get("test-prefix:foo3"));
+            assertEquals(new Tuple("bar4", 1234d), cluster.zpopmax("test-prefix:foo4"));
+        }
+    }
+}