DATAREDIS-696 - Add support for REPLACE option to RESTORE.

We bypass the lack of driver support by directly executing a custom command.

Original pull request: #344.

src/main/java/org/springframework/data/redis/connection/DefaultStringRedisConnection.java

@@ -1668,11 +1668,11 @@ public class DefaultStringRedisConnection implements StringRedisConnection, Deco
 
 	/*
 	 * (non-Javadoc)
-	 * @see org.springframework.data.redis.connection.RedisKeyCommands#restore(byte[], long, byte[])
+	 * @see org.springframework.data.redis.connection.RedisKeyCommands#restore(byte[], long, byte[], boolean)
 	 */
 	@Override
-	public void restore(byte[] key, long ttlInMillis, byte[] serializedValue) {
-		delegate.restore(key, ttlInMillis, serializedValue);
+	public void restore(byte[] key, long ttlInMillis, byte[] serializedValue, boolean replace) {
+		delegate.restore(key, ttlInMillis, serializedValue, replace);
 	}
 
 	/*

src/main/java/org/springframework/data/redis/connection/DefaultedRedisConnection.java

@@ -149,8 +149,8 @@ public interface DefaultedRedisConnection extends RedisConnection {
 	/** @deprecated in favor of {@link RedisConnection#keyCommands()}. */
 	@Override
 	@Deprecated
-	default void restore(byte[] key, long ttlInMillis, byte[] serializedValue) {
-		keyCommands().restore(key, ttlInMillis, serializedValue);
+	default void restore(byte[] key, long ttlInMillis, byte[] serializedValue, boolean replace) {
+		keyCommands().restore(key, ttlInMillis, serializedValue, replace);
 	}
 
 	/** @deprecated in favor of {@link RedisConnection#keyCommands()}. */

src/main/java/org/springframework/data/redis/connection/RedisKeyCommands.java

@@ -303,7 +303,21 @@ public interface RedisKeyCommands {
 	 * @param serializedValue must not be {@literal null}.
 	 * @see <a href="http://redis.io/commands/restore">Redis Documentation: RESTORE</a>
 	 */
-	void restore(byte[] key, long ttlInMillis, byte[] serializedValue);
+	default void restore(byte[] key, long ttlInMillis, byte[] serializedValue) {
+		restore(key, ttlInMillis, serializedValue, false);
+	}
+
+	/**
+	 * Create {@code key} using the {@code serializedValue}, previously obtained using {@link #dump(byte[])}.
+	 *
+	 * @param key must not be {@literal null}.
+	 * @param ttlInMillis
+	 * @param serializedValue must not be {@literal null}.
+	 * @param replace use {@literal true} to replace a potentially existing value instead of erroring.
+	 * @since 2.1
+	 * @see <a href="http://redis.io/commands/restore">Redis Documentation: RESTORE</a>
+	 */
+	void restore(byte[] key, long ttlInMillis, byte[] serializedValue, boolean replace);
 
 	/**
 	 * Get the type of internal representation used for storing the value at the given {@code key}.

src/main/java/org/springframework/data/redis/connection/jedis/JedisClusterKeyCommands.java

@@ -476,10 +476,10 @@ class JedisClusterKeyCommands implements RedisKeyCommands {
 
 	/*
 	 * (non-Javadoc)
-	 * @see org.springframework.data.redis.connection.RedisKeyCommands#restore(byte[], long, byte[])
+	 * @see org.springframework.data.redis.connection.RedisKeyCommands#restore(byte[], long, byte[], boolean)
 	 */
 	@Override
-	public void restore(byte[] key, long ttlInMillis, byte[] serializedValue) {
+	public void restore(byte[] key, long ttlInMillis, byte[] serializedValue, boolean replace) {
 
 		Assert.notNull(key, "Key must not be null!");
 		Assert.notNull(serializedValue, "Serialized value must not be null!");
@@ -488,9 +488,17 @@ class JedisClusterKeyCommands implements RedisKeyCommands {
 			throw new UnsupportedOperationException("Jedis does not support ttlInMillis exceeding Integer.MAX_VALUE.");
 		}
 
-		connection.getClusterCommandExecutor()
-				.executeCommandOnSingleNode((JedisClusterCommandCallback<String>) client -> client.restore(key,
-						Long.valueOf(ttlInMillis).intValue(), serializedValue), connection.clusterGetNodeForKey(key));
+		connection.getClusterCommandExecutor().executeCommandOnSingleNode((JedisClusterCommandCallback<String>) client -> {
+
+			if (!replace) {
+				return client.restore(key, Long.valueOf(ttlInMillis).intValue(), serializedValue);
+			}
+
+			return JedisConverters.toString(this.connection.execute("RESTORE", key,
+					Arrays.asList(JedisConverters.toBytes(ttlInMillis), serializedValue, JedisConverters.toBytes
+							("REPLACE"))));
+
+		}, connection.clusterGetNodeForKey(key));
 	}
 
 	/*

src/main/java/org/springframework/data/redis/connection/jedis/JedisKeyCommands.java

@@ -21,6 +21,7 @@ import redis.clients.jedis.ScanParams;
 import redis.clients.jedis.SortingParams;
 
 import java.time.Duration;
+import java.util.Arrays;
 import java.util.List;
 import java.util.Set;
 import java.util.concurrent.TimeUnit;
@@ -675,14 +676,22 @@ class JedisKeyCommands implements RedisKeyCommands {
 
 	/*
 	 * (non-Javadoc)
-	 * @see org.springframework.data.redis.connection.RedisKeyCommands#restore(byte[], long, byte[])
+	 * @see org.springframework.data.redis.connection.RedisKeyCommands#restore(byte[], long, byte[], boolean)
 	 */
 	@Override
-	public void restore(byte[] key, long ttlInMillis, byte[] serializedValue) {
+	public void restore(byte[] key, long ttlInMillis, byte[] serializedValue, boolean replace) {
 
 		Assert.notNull(key, "Key must not be null!");
 		Assert.notNull(serializedValue, "Serialized value must not be null!");
 
+		if (replace) {
+
+			this.connection.execute("RESTORE", new byte[][] { key, JedisConverters.toBytes(ttlInMillis),
+					serializedValue,
+					JedisConverters.toBytes("REPLACE") });
+			return;
+		}
+
 		if (ttlInMillis > Integer.MAX_VALUE) {
 			throw new IllegalArgumentException("TtlInMillis must be less than Integer.MAX_VALUE for restore in Jedis.");
 		}

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

@@ -21,6 +21,10 @@ import io.lettuce.core.ScanCursor;
 import io.lettuce.core.SortArgs;
 import io.lettuce.core.cluster.api.async.RedisClusterAsyncCommands;
 import io.lettuce.core.cluster.api.sync.RedisClusterCommands;
+import io.lettuce.core.codec.ByteArrayCodec;
+import io.lettuce.core.output.StatusOutput;
+import io.lettuce.core.protocol.CommandArgs;
+import io.lettuce.core.protocol.CommandType;
 import lombok.NonNull;
 import lombok.RequiredArgsConstructor;
 
@@ -671,23 +675,33 @@ class LettuceKeyCommands implements RedisKeyCommands {
 
 	/*
 	 * (non-Javadoc)
-	 * @see org.springframework.data.redis.connection.RedisKeyCommands#restore(byte[], long, byte[])
+	 * @see org.springframework.data.redis.connection.RedisKeyCommands#restore(byte[], long, byte[], boolean)
 	 */
 	@Override
-	public void restore(byte[] key, long ttlInMillis, byte[] serializedValue) {
+	public void restore(byte[] key, long ttlInMillis, byte[] serializedValue, boolean replace) {
 
 		Assert.notNull(key, "Key must not be null!");
 		Assert.notNull(serializedValue, "Serialized value must not be null!");
 
 		try {
+
+			if (replace) {
+
+				this.connection.execute("RESTORE", new byte[][] { key, LettuceConverters.toBytes(ttlInMillis), serializedValue,
+						LettuceConverters.toBytes("REPLACE") });
+				return;
+			}
+
 			if (isPipelined()) {
 				pipeline(connection.newLettuceStatusResult(getAsyncConnection().restore(key, ttlInMillis, serializedValue)));
 				return;
 			}
+
 			if (isQueueing()) {
 				transaction(connection.newLettuceStatusResult(getAsyncConnection().restore(key, ttlInMillis, serializedValue)));
 				return;
 			}
+
 			getConnection().restore(key, ttlInMillis, serializedValue);
 		} catch (Exception ex) {
 			throw convertLettuceAccessException(ex);

src/main/java/org/springframework/data/redis/core/RedisOperations.java

@@ -328,7 +328,23 @@ public interface RedisOperations<K, V> {
 	 * @param unit must not be {@literal null}.
 	 * @see <a href="http://redis.io/commands/restore">Redis Documentation: RESTORE</a>
 	 */
-	void restore(K key, byte[] value, long timeToLive, TimeUnit unit);
+	default void restore(K key, byte[] value, long timeToLive, TimeUnit unit) {
+		restore(key, value, timeToLive, unit, false);
+	}
+
+	/**
+	 * Create {@code key} using the {@code serializedValue}, previously obtained using {@link #dump(Object)}.
+	 *
+	 * @param key must not be {@literal null}.
+	 * @param value must not be {@literal null}.
+	 * @param timeToLive
+	 * @param unit must not be {@literal null}.
+	 * @param replace use {@literal true} to replace a potentially existing value instead of erroring.
+	 * @since 2.1
+	 * @see <a href="http://redis.io/commands/restore">Redis Documentation: RESTORE</a>
+	 */
+	void restore(K key, byte[] value, long timeToLive, TimeUnit unit, boolean replace);
+
 
 	/**
 	 * Get the time to live for {@code key} in seconds.

src/main/java/org/springframework/data/redis/core/RedisTemplate.java

@@ -980,16 +980,18 @@ public class RedisTemplate<K, V> extends RedisAccessor implements RedisOperation
 	 * @param value The value to restore, as returned by {@link #dump(Object)}
 	 * @param timeToLive An expiration for the restored key, or 0 for no expiration
 	 * @param unit The time unit for timeToLive
-	 * @throws RedisSystemException if the key you are attempting to restore already exists.
+	 * @param replace use {@literal true} to replace a potentially existing value instead of erroring.
+	 * @throws RedisSystemException if the key you are attempting to restore already exists and {@code replace} is set to
+	 *           {@literal false}.
 	 */
 	@Override
-	public void restore(K key, final byte[] value, long timeToLive, TimeUnit unit) {
+	public void restore(K key, final byte[] value, long timeToLive, TimeUnit unit, boolean replace) {
 
 		byte[] rawKey = rawKey(key);
 		long rawTimeout = TimeoutUtils.toMillis(timeToLive, unit);
 
 		execute(connection -> {
-			connection.restore(rawKey, rawTimeout, value);
+			connection.restore(rawKey, rawTimeout, value, replace);
 			return null;
 		}, true);
 	}

src/test/java/org/springframework/data/redis/connection/AbstractConnectionIntegrationTests.java

@@ -1113,6 +1113,20 @@ public abstract class AbstractConnectionIntegrationTests {
 		getResults();
 	}
 
+	@Test // DATAREDIS-696
+	@IfProfileValue(name = "redisVersion", value = "3.0+")
+	public void testRestoreExistingKeyWithReplaceOption() {
+
+		actual.add(connection.set("testing", "12"));
+		actual.add(connection.dump("testing".getBytes()));
+		actual.add(connection.set("testing", "21"));
+		connection.restore("testing".getBytes(), 0, (byte[]) getResults().get(1), true);
+
+		initConnection();
+		actual.add(connection.get("testing"));
+		verifyResults(Arrays.asList(new Object[] { "12" }));
+	}
+
 	@Test
 	@IfProfileValue(name = "redisVersion", value = "2.6+")
 	public void testRestoreTtl() {
@@ -2947,14 +2961,10 @@ public abstract class AbstractConnectionIntegrationTests {
 	@WithRedisDriver({ RedisDriver.JEDIS, RedisDriver.LETTUCE })
 	public void bitFieldIncrByWithOverflowShouldWorkCorrectly() {
 
-		actual.add(
-				connection.bitfield(KEY_1, create().incr(unsigned(2)).valueAt(offset(102L)).overflow(FAIL).by(1L)));
-		actual.add(
-				connection.bitfield(KEY_1, create().incr(unsigned(2)).valueAt(offset(102L)).overflow(FAIL).by(1L)));
-		actual.add(
-				connection.bitfield(KEY_1, create().incr(unsigned(2)).valueAt(offset(102L)).overflow(FAIL).by(1L)));
-		actual.add(
-				connection.bitfield(KEY_1, create().incr(unsigned(2)).valueAt(offset(102L)).overflow(FAIL).by(1L)));
+		actual.add(connection.bitfield(KEY_1, create().incr(unsigned(2)).valueAt(offset(102L)).overflow(FAIL).by(1L)));
+		actual.add(connection.bitfield(KEY_1, create().incr(unsigned(2)).valueAt(offset(102L)).overflow(FAIL).by(1L)));
+		actual.add(connection.bitfield(KEY_1, create().incr(unsigned(2)).valueAt(offset(102L)).overflow(FAIL).by(1L)));
+		actual.add(connection.bitfield(KEY_1, create().incr(unsigned(2)).valueAt(offset(102L)).overflow(FAIL).by(1L)));
 
 		List<Object> results = getResults();
 		assertThat((List<Long>) results.get(0), contains(1L));
@@ -2969,8 +2979,7 @@ public abstract class AbstractConnectionIntegrationTests {
 	public void bitfieldShouldAllowMultipleSubcommands() {
 
 		actual.add(
-				connection.bitfield(KEY_1,
-				create().incr(signed(5)).valueAt(offset(100L)).by(1L).get(unsigned(4)).valueAt(0L)));
+				connection.bitfield(KEY_1, create().incr(signed(5)).valueAt(offset(100L)).by(1L).get(unsigned(4)).valueAt(0L)));
 
 		assertThat((List<Long>) getResults().get(0), contains(1L, 0L));
 	}

src/test/java/org/springframework/data/redis/connection/ClusterConnectionTests.java

@@ -96,6 +96,9 @@ public interface ClusterConnectionTests {
 	// DATAREDIS-315
 	void dumpAndRestoreShouldWorkCorrectly();
 
+	// DATAREDIS-696
+	void dumpAndRestoreWithReplaceOptionShouldWorkCorrectly();
+
 	// DATAREDIS-315
 	void echoShouldReturnInputCorrectly();
 

src/test/java/org/springframework/data/redis/connection/RedisConnectionUnitTests.java

@@ -772,8 +772,8 @@ public class RedisConnectionUnitTests {
 			delegate.slaveOfNoOne();
 		}
 
-		public void restore(byte[] key, long ttlInMillis, byte[] serializedValue) {
-			delegate.restore(key, ttlInMillis, serializedValue);
+		public void restore(byte[] key, long ttlInMillis, byte[] serializedValue, boolean replace) {
+			delegate.restore(key, ttlInMillis, serializedValue, replace);
 		}
 
 		public byte[] bRPopLPush(int timeout, byte[] srcKey, byte[] dstKey) {

src/test/java/org/springframework/data/redis/connection/jedis/JedisClusterConnectionTests.java

@@ -355,6 +355,20 @@ public class JedisClusterConnectionTests implements ClusterConnectionTests {
 		assertThat(nativeConnection.get(KEY_2), is(VALUE_1));
 	}
 
+	@Test // DATAREDIS-696
+	public void dumpAndRestoreWithReplaceOptionShouldWorkCorrectly() {
+
+		nativeConnection.set(KEY_1, VALUE_1);
+
+		byte[] dumpedValue = clusterConnection.keyCommands().dump(KEY_1_BYTES);
+
+		nativeConnection.set(KEY_1, VALUE_2);
+
+		clusterConnection.keyCommands().restore(KEY_1_BYTES, 0, dumpedValue, true);
+
+		assertThat(nativeConnection.get(KEY_1), is(VALUE_1));
+	}
+
 	@Test // DATAREDIS-315
 	public void echoShouldReturnInputCorrectly() {
 		assertThat(clusterConnection.echo(VALUE_1_BYTES), is(VALUE_1_BYTES));

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

@@ -392,6 +392,20 @@ public class LettuceClusterConnectionTests implements ClusterConnectionTests {
 		assertThat(nativeConnection.get(KEY_2), is(VALUE_1));
 	}
 
+	@Test // DATAREDIS-696
+	public void dumpAndRestoreWithReplaceOptionShouldWorkCorrectly() {
+
+		nativeConnection.set(KEY_1, VALUE_1);
+
+		byte[] dumpedValue = clusterConnection.keyCommands().dump(KEY_1_BYTES);
+
+		nativeConnection.set(KEY_1, VALUE_2);
+
+		clusterConnection.keyCommands().restore(KEY_1_BYTES, 0, dumpedValue, true);
+
+		assertThat(nativeConnection.get(KEY_1), is(VALUE_1));
+	}
+
 	@Test // DATAREDIS-315
 	public void echoShouldReturnInputCorrectly() {
 		assertThat(clusterConnection.echo(VALUE_1_BYTES), is(VALUE_1_BYTES));