Ricky Hao

Python:类内装饰器的使用(pymongo自动重连实现)

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装饰即可。

点赞

发表评论

电子邮件地址不会被公开。 必填项已用*标注

此站点使用Akismet来减少垃圾评论。了解我们如何处理您的评论数据