import redis handle dirs

scripts/import_redis_db.py

@@ -5,7 +5,7 @@
 
 parser = argparse.ArgumentParser(description='Export a DB from Redis')
 parser.add_argument('--db', type=str, required=True,help='DB to populate')
-parser.add_argument('--csv', type=str, required=True,help='Input file in CSV format')
+parser.add_argument('--csv', type=str, required=True,help='Input directory containing CSV files')
 parser.add_argument('--key', type=str, required=True,help='column (i.e. number of column) to use as key in redis')
 parser.add_argument('--value', type=str, required=True,help='column (i.e. number of column) to use as value in redis')
 
@@ -30,20 +30,29 @@
 
 r_db = Redis(host="localhost", port="6379", db=args.db)
 
-print("Insert in redis DB="+str(args.db)+", the data in="+str(args.csv)+" .")
-with open(args.csv,'r') as f:
-    reader = csv.reader(f)
-    w_buffer = dict()
-    REDIS_W_BUFFER = 10000
-    for row in tqdm(reader):
-        key = row[int(args.key)]
-        value = row[int(args.value)]
-        w_buffer[key] = value
-        if len(w_buffer.keys()) >= REDIS_W_BUFFER:
-            r_db.mset(w_buffer)
+print("Insert in redis DB="+str(args.db)+", the data inside the CSVs in the directoy="+str(args.csv)+" .")
+for filename in os.listdir(args.csv):
+    if filename.endswith('.csv'):
+        with open(filename,'r') as f:
+            reader = csv.reader(f)
             w_buffer = dict()
-
-    if len(w_buffer.keys()) > 0:
-        r_db.mset(w_buffer)
+            REDIS_W_BUFFER = 10000
+            for row in tqdm(reader):
+                key = row[int(args.key)]
+                value = row[int(args.value)]
+                w_buffer[key] = value
+                if len(w_buffer.keys()) >= REDIS_W_BUFFER:
+                    r_db.mset(w_buffer)
+                    w_buffer = dict()
+
+            if len(w_buffer.keys()) > 0:
+                r_db.mset(w_buffer)
 
 print("Done!")
+
+
+
+csv_files = []
+    for filename in os.listdir(directory):
+        if filename.endswith('.csv'):
+            csv_files.append(os.path.join(directory, filename))