0x00

最近要实现一个pymongo库的自动重连功能，但是官方文档只说了：

In order to auto-reconnect you must handle this exception, recognizing that the operation which caused it has not necessarily succeeded. Future operations will attempt to open a new connection to the database (and will continue to raise this exception until the first successful connection is made).

大概意思就是。。想要自动重连？自己捕获异常重连吧~
因此，这里简易地用装饰器来实现一个自动重连的功能。

0x01

重连函数

def reconnect(func):
    """ Reconnect decorator. """
    @wraps(func)
    def wrapper(self, *args, **kwargs):
        """ Wrapper. """
        try:
            return func(self, *args, **kwargs)
        except errors.AutoReconnect:
            logger.warning('[W] Reconnecting to database.')
            #self.connect()
            return wrapper(self, *args, **kwargs)
    return wrapper

这里采用装饰器的方式，当捕获到函数运行时的AutoReconnect错误时，就自动重连。
因为我采用的是类封装，因此传入了self参数。（这个函数不能定义在类里，定义在类外面。）

0x02

封装的数据库类

class Database:
    """
    Database class
    """
    def __init__(self):
        self.connect()

    def connect(self):
        """ Database connect function. """
        self.client = MongoClient(host=config['DATABASE_HOST'], 
                authSource=config['DATABASE_NAME'], 
                username=config['DATABASE_USER'], 
                password=config['DATABASE_PASSWORD'])
        logger.info("[+] Database connected.")

    @reconnect
    def insert_one(self, database, collection, document):
        """ Insert one document into specific database collection. """
        return self.client[database][collection].insert_one(document).inserted_id

大概就像这样，在可能需要重连的操作函数上加上@reconnect装饰即可。