BGP-Ranking/tools/migrate.py

69 lines
2.1 KiB
Python

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
from datetime import datetime
from typing import Set
from redis import Redis
redis_src = Redis(unix_socket_path='../storage/storage.sock')
redis_dst = Redis('127.0.0.1', 5188)
chunk_size = 100000
def process_chunk(src: Redis, dst: Redis, keys: Set[str]):
src_pipeline = src.pipeline()
[src_pipeline.type(key) for key in keys]
to_process = {key: key_type for key, key_type in zip(keys, src_pipeline.execute())}
src_pipeline = src.pipeline()
for key, key_type in to_process.items():
if key_type == b"string":
src_pipeline.get(key)
elif key_type == b"list":
raise Exception('Lists should not be used.')
elif key_type == b"set":
src_pipeline.smembers(key)
elif key_type == b"zset":
src_pipeline.zrangebyscore(key, '-Inf', '+Inf', withscores=True)
elif key_type == b"hash":
src_pipeline.hgetall(key)
else:
raise Exception(f'{key_type} not supported {key}.')
dest_pipeline = dst.pipeline()
for key, content in zip(to_process.keys(), src_pipeline.execute()):
if to_process[key] == b"string":
dest_pipeline.set(key, content)
elif to_process[key] == b"set":
dest_pipeline.sadd(key, *content)
elif to_process[key] == b"zset":
dest_pipeline.zadd(key, {value: rank for value, rank in content})
elif to_process[key] == b"hash":
dest_pipeline.hmset(key, content)
dest_pipeline.execute()
def migrate(src: Redis, dst: Redis):
keys = set()
pos = 0
for key in src.scan_iter(count=chunk_size, match='2017*'):
keys.add(key)
if len(keys) == chunk_size:
process_chunk(src, dst, keys)
pos += len(keys)
print(f'{datetime.now()} - {pos} keys done.')
keys = set()
# migrate remaining keys
process_chunk(src, dst, keys)
pos += len(keys)
print(f'{datetime.now()} - {pos} keys done.')
if __name__ == '__main__':
migrate(redis_src, redis_dst)