From dfeecba0411a4517abb278d2fe3bd032d7d57ae8 Mon Sep 17 00:00:00 2001
From: antirez <antirez@gmail.com>
Date: Wed, 7 Nov 2012 10:57:23 +0100
Subject: [PATCH] REPLACE option for RESTORE.

The REPLACE option deletes an existing key with the same name (if any)
and materializes the new one. The default behavior without RESTORE is to
return an error if a key already exists.
---
 src/cluster.c       | 19 ++++++++++++++++---
 src/redis.c         |  2 +-
 tests/unit/dump.tcl | 16 ++++++++++++++++
 3 files changed, 33 insertions(+), 4 deletions(-)

diff --git a/src/cluster.c b/src/cluster.c
index 57243132d..33a267589 100644
--- a/src/cluster.c
+++ b/src/cluster.c
@@ -1539,15 +1539,25 @@ void dumpCommand(redisClient *c) {
     return;
 }
 
-/* RESTORE key ttl serialized-value */
+/* RESTORE key ttl serialized-value [REPLACE] */
 void restoreCommand(redisClient *c) {
     long ttl;
     rio payload;
-    int type;
+    int j, type, replace = 0;
     robj *obj;
 
+    /* Parse additional options */
+    for (j = 4; j < c->argc; j++) {
+        if (!strcasecmp(c->argv[j]->ptr,"replace")) {
+            replace = 1;
+        } else {
+            addReply(c,shared.syntaxerr);
+            return;
+        }
+    }
+
     /* Make sure this key does not already exist here... */
-    if (lookupKeyWrite(c->db,c->argv[1]) != NULL) {
+    if (!replace && lookupKeyWrite(c->db,c->argv[1]) != NULL) {
         addReplyError(c,"Target key name is busy.");
         return;
     }
@@ -1574,6 +1584,9 @@ void restoreCommand(redisClient *c) {
         return;
     }
 
+    /* Remove the old key if needed. */
+    if (replace) dbDelete(c->db,c->argv[1]);
+
     /* Create the key and set the TTL if any */
     dbAdd(c->db,c->argv[1],obj);
     if (ttl) setExpire(c->db,c->argv[1],mstime()+ttl);
diff --git a/src/redis.c b/src/redis.c
index 5ac025e75..0780e8876 100644
--- a/src/redis.c
+++ b/src/redis.c
@@ -240,7 +240,7 @@ struct redisCommand redisCommandTable[] = {
     {"watch",watchCommand,-2,"rs",0,noPreloadGetKeys,1,-1,1,0,0},
     {"unwatch",unwatchCommand,1,"rs",0,NULL,0,0,0,0,0},
     {"cluster",clusterCommand,-2,"ar",0,NULL,0,0,0,0,0},
-    {"restore",restoreCommand,4,"awm",0,NULL,1,1,1,0,0},
+    {"restore",restoreCommand,-4,"awm",0,NULL,1,1,1,0,0},
     {"migrate",migrateCommand,6,"aw",0,NULL,0,0,0,0,0},
     {"asking",askingCommand,1,"r",0,NULL,0,0,0,0,0},
     {"dump",dumpCommand,2,"ar",0,NULL,1,1,1,0,0},
diff --git a/tests/unit/dump.tcl b/tests/unit/dump.tcl
index be891a96f..1eb91eb21 100644
--- a/tests/unit/dump.tcl
+++ b/tests/unit/dump.tcl
@@ -23,6 +23,22 @@ start_server {tags {"dump"}} {
         set e
     } {*is busy*}
 
+    test {RESTORE can overwrite an existing key with REPLACE} {
+        r set foo bar1
+        set encoded1 [r dump foo]
+        r set foo bar2
+        set encoded2 [r dump foo]
+        r del foo
+        r restore foo 0 $encoded1
+        r restore foo 0 $encoded2 replace
+        r get foo
+    } {bar2}
+
+    test {RESTORE can detect a syntax error for unrecongized options} {
+        catch {r restore foo 0 "..." invalid-option} e
+        set e
+    } {*syntax*}
+
     test {DUMP of non existing key returns nil} {
         r dump nonexisting_key
     } {}