From 541190be4dfbf8a183a390540cb459aedfb6df78 Mon Sep 17 00:00:00 2001 From: ivanhb Date: Mon, 25 Mar 2024 14:01:40 +0100 Subject: [PATCH] import redis handle dirs --- scripts/import_redis_db.py | 39 +++++++++++++++++++++++--------------- 1 file changed, 24 insertions(+), 15 deletions(-) diff --git a/scripts/import_redis_db.py b/scripts/import_redis_db.py index 4db7c70..f879f71 100755 --- a/scripts/import_redis_db.py +++ b/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))