Added Reds Cluster support to redis-copy script

README.markdown

@@ -10,14 +10,16 @@ It accepts as input a pipe with redis commands formatted as "DEL key", "SET key
 
 
 #**Redis Copy**
+NOW SUPPORTS REDIS-CLUSTER DESTINATIONS
 
 Redis Copy the keys in a source redis server into another target redis server.
 The script probably needs to be added to a cron job if the keys are a lot because it only copies a fix number of keys at a time
 and continue from there on the next run. It does this until there is no more keys to copy
 
 ####Dependency:
 
-    sudo easy_install -U redis
+    sudo easy_install -U redis  (OR: sudo pip install redis)
+    redis-py-cluster  (clone https://github.com/Grokzen/redis-py-cluster and run: sudo python setup.py install)
 
 ####Usage:
 
@@ -30,6 +32,9 @@ and continue from there on the next run. It does this until there is no more key
     -t ..., --target=...        target redis server "ip:port" to copy keys to. e.g. 192.168.0.101:6379
     -d ..., --databases=...     comma separated list of redis databases to select when copying. e.g. 2,5
     -h, --help                  show this help
+    -S ..., --prefix=...        optional to prefix destination key with value (this is overwritten if --clustered=2)
+    -T ..., --totable=...       optional to target a different destination database table (not used if dest is a cluster)
+    -C ..., --clustered=...     optional 0, 1, 2. Use 1 if the destination is a cluster. Use 2 to prefix the origin DB name into the destination key. Ex: "0_keyname"
     --clean                     clean all variables, temp lists created previously by the script
 
 
@@ -54,6 +59,14 @@ and continue from there on the next run. It does this until there is no more key
     --databases=2,5                                         copy all keys in db 2 and 5 from server 192.168.0.99:6379 to server 192.168.0.101:6379
                                                           with a limit of 1000 per script run
 
+    python redis-copy.py \
+    --source=192.168.0.99:6379 \
+    --target=192.168.0.101:6379 \
+    --databases=1 \
+    --clustered=2                                       copy from non-clustered db of table 1 into the destination cluster. The dest table will be 0 
+                                                        as defined by redis-cluster. Since --clustered=2, each dest key will be 
+                                                        prepended with "DB_keyname". In this case, keys will be named "1_keyname".
+
 
 
 

redis-copy.py

@@ -15,9 +15,14 @@
   -t ..., --target=...        target redis server "ip:port" to copy keys to. e.g. 192.168.0.101:6379
   -d ..., --databases=...     comma separated list of redis databases to select when copying. e.g. 2,5
   -h, --help                  show this help
+  -S ..., --prefix=...        optional to prefix destination key with value (this is overwritten if --clustered=2)
+  -T ..., --totable=...       optional to target a different destination database table  (not used if dest is a cluster)
+  -C ..., --clustered=...     optional 0, 1, 2. Use 1 if the destination is a cluster. Use 2 to prefix the origin DB name into the destination key. Ex: "0_keyname"
   --clean                     clean all variables, temp lists created previously by the script
 
-Dependencies: redis (redis-py: sudo pip install redis)
+**Dependencies:** 
+redis (redis-py: sudo pip install redis)
+redis-py-cluster (clone https://github.com/Grokzen/redis-py-cluster and run: sudo python setup.py install)
 
 Examples:
   python redis-copy.py --help                                show this doc
@@ -39,6 +44,14 @@
   --databases=2,5                                         copy all keys in db 2 and 5 from server 192.168.0.99:6379 to server 192.168.0.101:6379
                                                           with a limit of 1000 per script run
 
+  python redis-copy.py \
+  --source=192.168.0.99:6379 \
+  --target=192.168.0.101:6379 \
+  --databases=1 \
+  --clustered=2 \                                       copy from non-clustered db of table 1 into the destination cluster. The dest table will be 0 
+                                                        as defined by redis-cluster. Since --clustered=2, each dest key will be 
+                                                        prepended with "DB_keyname". In this case, keys will be named "1_keyname".
+
 """
 
 __author__ = "Salimane Adjao Moustapha ([email protected])"
@@ -52,6 +65,7 @@
 import time
 import sys
 import getopt
+from rediscluster import StrictRedisCluster
 
 
 class RedisCopy:
@@ -66,10 +80,13 @@ class RedisCopy:
     # numbers of keys to copy on each iteration
     limit = 10000
 
-    def __init__(self, source, target, dbs):
+    def __init__(self, source, target, dbs, prefix, toTable, clustered):
         self.source = source
         self.target = target
         self.dbs = dbs
+        self.prefix = prefix
+        self.toTable = toTable;
+        self.clustered = clustered;
 
     def save_keylists(self):
         """Function to save the keys' names of the source redis server into a list for later usage.
@@ -134,13 +151,21 @@ def copy_db(self, limit=None):
             print "Started copy of %s keys from %d to %d at %s...\n" % (servername, keymoved, dbsize, time.strftime("%Y-%m-%d %I:%M:%S"))
 
             #get redis handle for corresponding target server-db
-            rr = redis.StrictRedis(
-                host=self.target['host'], port=self.target['port'], db=db)
+            destTable = db if self.toTable==-1 else self.toTable;
+
+            if self.clustered=="0":
+                rr = redis.StrictRedis(host=self.target['host'], port=self.target['port'], db=destTable)
+            else:
+                rr = StrictRedisCluster(startup_nodes=[{'host':self.target['host'], 'port':self.target['port']}])
+                if self.clustered=="2":
+                    self.prefix = str(db) + "_"
+                    print "Using table prefix of: %s" % (self.prefix)
 
             #max index for lrange
             newkeymoved = keymoved + \
                 self.limit if dbsize > keymoved + self.limit else dbsize
 
+            prefix = self.prefix
             for key in r.lrange(self.mprefix + self.keylistprefix + servername, keymoved, newkeymoved):
                 #get key type
                 ktype = r.type(key)
@@ -150,32 +175,32 @@ def copy_db(self, limit=None):
 
                 #save key to target server-db
                 if ktype == 'string':
-                    rr.set(key, r.get(key))
+                    rr.set(prefix+key, r.get(key))
                 elif ktype == 'hash':
-                    rr.hmset(key, r.hgetall(key))
+                    rr.hmset(prefix+key, r.hgetall(key))
                 elif ktype == 'list':
                     if key == self.mprefix + "keylist:" + servername:
                         continue
                     #value = r.lrange(key, 0, -1)
                     #rr.rpush(key, *value)
                     for k in r.lrange(key, 0, -1):
-                        rr.rpush(key, k)
+                        rr.rpush(prefix+key, k)
                 elif ktype == 'set':
                     #value = r.smembers(key)
                     #rr.sadd(key, *value)
                     for k in r.smembers(key):
-                        rr.sadd(key, k)
+                        rr.sadd(prefix+key, k)
                 elif ktype == 'zset':
                     #value = r.zrange(key, 0, -1, withscores=True)
                     #rr.zadd(key, **dict(value))
                     for k, v in r.zrange(key, 0, -1, withscores=True):
-                        rr.zadd(key, v, k)
+                        rr.zadd(prefix+key, v, k)
 
                 # Handle keys with an expire time set
                 kttl = r.ttl(key)
                 kttl = -1 if kttl is None else int(kttl)
                 if kttl != -1:
-                    rr.expire(key, kttl)
+                    rr.expire(prefix+key, kttl)
 
                 moved += 1
 
@@ -191,12 +216,18 @@ def flush_target(self):
         """Function to flush the target server.
         """
         for db in self.dbs:
+            destTable = db if self.toTable==-1 else self.toTable;
+
             servername = self.target['host'] + ":" + str(
-                self.target['port']) + ":" + str(db)
+                self.target['port']) + ":" + str(destTable)
             print "Flushing server %s at %s...\n" % (
                 servername, time.strftime("%Y-%m-%d %I:%M:%S"))
-            r = redis.StrictRedis(
-                host=self.target['host'], port=self.target['port'], db=db)
+
+            if self.clustered==0:
+                r = redis.StrictRedis(host=self.target['host'], port=self.target['port'], db=destTable)
+            else:
+                r = StrictRedisCluster(startup_nodes=[{'host':self.target['host'], 'port':self.target['port']}])
+
             r.flushdb()
             print "Flushed server %s at %s...\n" % (
                 servername, time.strftime("%Y-%m-%d %I:%M:%S"))
@@ -219,7 +250,7 @@ def clean(self):
         print "Done.\n"
 
 
-def main(source, target, databases, limit=None, clean=False):
+def main(source, target, databases, limit=None, clean=False, prefix="", toTable="-1", clustered="0"):
     #getting source and target
     if (source == target):
         exit('The 2 servers adresses are the same. e.g. python redis-copy.py 127.0.0.1:6379 127.0.0.1:63791  0,1')
@@ -246,7 +277,7 @@ def main(source, target, databases, limit=None, clean=False):
     except AttributeError as e:
         exit('Please this script requires redis-py >= 2.4.10, your current version is :' + redis.__version__)
 
-    mig = RedisCopy(source_server, target_server, dbs)
+    mig = RedisCopy(source_server, target_server, dbs, prefix, toTable, clustered)
 
     if clean == False:
         #check if script already running
@@ -276,8 +307,11 @@ def usage():
 
 if __name__ == "__main__":
     clean = False
+    prefix = "";
+    toTable = "-1";
+    clustered = "0";
     try:
-        opts, args = getopt.getopt(sys.argv[1:], "hl:s:t:d:", ["help", "limit=", "source=", "target=", "databases=", "clean"])
+        opts, args = getopt.getopt(sys.argv[1:], "hl:s:t:d:S:T:C:", ["help", "limit=", "source=", "target=", "databases=", "clean", "prefix=", "totable=", "clustered="])
     except getopt.GetoptError:
         usage()
         sys.exit(2)
@@ -295,13 +329,20 @@ def usage():
             target = arg
         elif opt in ("-d", "--databases"):
             databases = arg
+        elif opt in ("-S", "--prefix"):
+            prefix = arg
+        elif opt in ("-T", "--totable"):
+            toTable = arg
+        elif opt in ("-C", "--clustered"):
+            clustered = arg
+
 
     try:
         limit = int(limit)
     except (NameError, TypeError, ValueError):
         limit = None
 
     try:
-        main(source, target, databases, limit, clean)
+        main(source, target, databases, limit, clean, prefix, toTable, clustered)
     except NameError as e:
         usage()