python实现redis rdb迁移

前言

找了一圈redis迁移工具，包括：redis-port、redis-shake、redis-dump等等，涉及到各种语言和环境，折腾了半天没有成功。。最后还是自己写吧。

利用https://github.com/sripathikrishnan/redis-rdb-tools这个强大的python库简单写一个适用于单机的redis迁移脚本。

python脚本

import datetime
from typing import Dict, Optional, List

import redis
from loguru import logger
from rdbtools import RdbParser, RdbCallback
from rdbtools.encodehelpers import bytes_to_unicode, STRING_ESCAPE_UTF8


class RedisMigrator:
    """redis迁移"""
    def __init__(self, rdb_file: str, redis_config: Dict, source_dbs: Optional[List[int]] = None):
        self.rdb_file = rdb_file
        self.redis_config = redis_config
        self.source_dbs = source_dbs
        self.redis_client = self.init_redis()

    def init_redis(self):
        """初始化redis客户端"""
        redis_pool = redis.ConnectionPool(**self.redis_config)
        redis_client = redis.Redis(connection_pool=redis_pool)
        return redis_client

    def migrate(self):
        """
        迁移
        filters参数支持多种条件过滤
        {"dbs" : [0, 1], "keys" : "foo.*", "types" : ["hash", "set", "sortedset", "list", "string"]}
        """
        callback = MyCallback(redis_client=self.redis_client)
        filters = None
        if self.source_dbs:
            filters = {"dbs": self.source_dbs}
        parser = RdbParser(callback, filters=filters)
        parser.parse(self.rdb_file)


class MyCallback(RdbCallback):
    """回调"""
    def __init__(self, redis_client):
        super(MyCallback, self).__init__(string_escape=STRING_ESCAPE_UTF8)
        self.redis_client = redis_client

    def encode_key(self, key):
        return bytes_to_unicode(key, self._escape, skip_printable=True)

    def encode_value(self, val):
        return bytes_to_unicode(val, self._escape)

    def set(self, key, value, expiry, info):
        logger.info("set")
        key = self.encode_key(key)
        value = self.encode_value(value)
        if expiry:
            now = datetime.datetime.now()
            # 计算距离当前时间的秒数
            expiry = int((expiry - now).total_seconds())
        logger.info("{} {} {}".format(key, value, expiry))

        if not expiry:
            self.redis_client.set(key, value)
        elif expiry > 0:
            self.redis_client.setex(key, expiry, value)

    def hset(self, key, field, value):
        logger.info("hset")
        key = self.encode_key(key)
        value = self.encode_value(value)
        field = self.encode_value(field)
        logger.info("{} {}".format(key, value))

        self.redis_client.hset(key, field, value)

    def rpush(self, key, value):
        logger.info("rpush")
        key = self.encode_key(key)
        value = self.encode_value(value)
        logger.info("{} {}".format(key, value))

        self.redis_client.rpush(key, value)

    def zadd(self, key, score, member):
        logger.info("zadd")
        key = self.encode_key(key)
        score = self.encode_value(score)
        member = self.encode_value(member)
        logger.info("{} {} {}".format(key, score, member))

        self.redis_client.zadd(key, {member: score})

    def sadd(self, key, member):
        logger.info("sadd")
        key = self.encode_key(key)
        value = self.encode_value(member)
        logger.info("{} {}".format(key, value))

        self.redis_client.sadd(key, value)

    def start_database(self, db_number):
        logger.info("Starting database {}".format(db_number))

    def end_database(self, db_number):
        logger.info("End database {}".format(db_number))

if __name__ == '__main__':
    # 目标redis配置
    target_redis_config = {
        "host": "127.0.0.1",
        "port": 6379,
        "db": 13,
        "password": "123456"
    }
    # 要迁移的rdb文件
    _rdb_file = r'./dump.rdb'
    # 要迁移的数据库列表，不传代表全部0-15
    _source_dbs = [13]

    migrator = RedisMigrator(rdb_file=_rdb_file, redis_config=target_redis_config, source_dbs=_source_dbs)
    migrator.migrate()