From f560c0533fcf5cffec12761caa0a9322921c5cc4 Mon Sep 17 00:00:00 2001
From: holys <chendahui007@gmail.com>
Date: Wed, 8 Oct 2014 11:33:43 +0800
Subject: [PATCH] add data migration support for data type Set

---
 tools/redis_import/README.md       |  1 -
 tools/redis_import/redis_import.py | 12 +++++++++++-
 tools/redis_import/test.py         | 12 +++++++++---
 3 files changed, 20 insertions(+), 5 deletions(-)

diff --git a/tools/redis_import/README.md b/tools/redis_import/README.md
index 634ed60..e93f847 100644
--- a/tools/redis_import/README.md
+++ b/tools/redis_import/README.md
@@ -1,6 +1,5 @@
 ## Notice
 
-1. The tool doesn't support `set` data type.
 2. The tool doesn't support `bitmap` data type.
 2. Our `zset` use integer instead of double, so the zset float score in Redis 
    will be **converted to integer**.
diff --git a/tools/redis_import/redis_import.py b/tools/redis_import/redis_import.py
index 7a463ae..7036f89 100644
--- a/tools/redis_import/redis_import.py
+++ b/tools/redis_import/redis_import.py
@@ -35,7 +35,7 @@ def set_ttl(redis_client, ledis_client, key, k_type):
         "string": ledis_client.expire,
         "list": ledis_client.lexpire,
         "hash": ledis_client.hexpire,
-        "set": ledis_client.zexpire,
+        "set": ledis_client.sexpire,
         "zset": ledis_client.zexpire
     }
     timeout = redis_client.ttl(key)
@@ -46,6 +46,7 @@ def set_ttl(redis_client, ledis_client, key, k_type):
 def copy_key(redis_client, ledis_client, key, convert=False):
     global entries
     k_type = redis_client.type(key)
+
     if k_type == "string":
         value = redis_client.get(key)
         ledis_client.set(key, value)
@@ -66,6 +67,7 @@ def copy_key(redis_client, ledis_client, key, convert=False):
         entries += 1
 
     elif k_type == "zset":
+        # dangerous to do this?
         out = redis_client.zrange(key, 0, -1, withscores=True)
         pieces = od()
         for i in od(out).iteritems():
@@ -74,6 +76,14 @@ def copy_key(redis_client, ledis_client, key, convert=False):
         set_ttl(redis_client, ledis_client, key, k_type)
         entries += 1
 
+    elif k_type == "set":
+        mbs = list(redis_client.smembers(key))
+
+        if mbs is not None:
+            ledis_client.sadd(key, *mbs)
+            set_ttl(redis_client, ledis_client, key, k_type)
+            entries += 1
+
     else:
         print "KEY %s of TYPE %s is not supported by LedisDB." % (key, k_type)
 
diff --git a/tools/redis_import/test.py b/tools/redis_import/test.py
index 96cceeb..0279c26 100644
--- a/tools/redis_import/test.py
+++ b/tools/redis_import/test.py
@@ -47,6 +47,9 @@ def random_zset(client, words, length=1000):
     client.zadd("zsetName", **d)
 
 
+def random_set(client, words, length=1000):
+    client.sadd("setName", *words)
+
 def test():
     words = get_words()
     print "Flush all redis data before insert new."
@@ -64,10 +67,9 @@ def test():
     random_zset(rds, words)
     print "random_zset done"
 
+    random_set(rds, words)
+    print "random_set done"
 
-    lds.lclear("listName")
-    lds.hclear("hashName")
-    lds.zclear("zsetName")
     copy(rds, lds, convert=True)
 
     # for all keys
@@ -101,6 +103,7 @@ def ledis_ttl(ledis_client, key, k_type):
         "list": lds.lttl,
         "hash": lds.httl,
         "zset": lds.zttl,
+        "set": lds.sttl,
     }
     return ttls[k_type](key)
 
@@ -115,6 +118,9 @@ def test_ttl():
             assert ledis_ttl(lds, key, k_type) > 0
 
 if __name__ == "__main__":
+    rds.flushdb()
+    lds.flushdb()
+
     test()
     test_ttl()
     print "Test passed."