AutoReconnectingConnection

connection.py

import functools
import logging
from pymongo.errors import AutoReconnect
from pymongo import Connection
from pymongo import ReplicaSetConnection
import time
import types

def wrap(func, num_retries, sleep_interval):
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        trial = 0
        while trial <= num_retries:
            try:
                return func(*args, **kwargs)
            except AutoReconnect, e:
                trial += 1
                if trial > num_retries:
                    raise e
                logging.info('AutoReconnect raised, msg %s.. retrying..' % str(e))
                if sleep_interval:
                    time.sleep(sleep_interval)
    return wrapper

class AutoReconnectingConnection(Connection):
    def __init__(self, *args, **kwargs):
        """
        @param connection: pymongo.Connection, connection to wrap
        @param num_retries: int, how many times to retry the same operation
        @param sleep_interval: int, how long to sleep in between retries, if
        any.
        """
        self._num_retries = kwargs.get('num_retries', 1)
        self._sleep_interval = kwargs.get('sleep_interval', None)
        del(kwargs['sleep_interval'])
        del(kwargs['num_retries'])
        wrap(Connection.__init__, self._num_retries, self._sleep_interval)(self, *args, **kwargs)

    def __getattr__(self, name):
        func = Connection.__getattr__(self, name)

        if isinstance(func, types.FunctionType):
            return wrap(func, self._num_retries, self._sleep_interval)
        else:
            return func

    def __getitem__(self, name):
        return self.__getattr__(name)

class AutoReconnectingReplicaSetConnection(ReplicaSetConnection):
    def __init__(self, *args, **kwargs):
        """
        @param connection: pymongo.Connection, connection to wrap
        @param num_retries: int, how many times to retry the same operation
        @param sleep_interval: int, how long to sleep in between retries, if
        any.
        """
        self._num_retries = kwargs.get('num_retries', 1)
        self._sleep_interval = kwargs.get('sleep_interval', None)
        del(kwargs['sleep_interval'])
        del(kwargs['num_retries'])
        wrap(ReplicaSetConnection.__init__, self._num_retries, self._sleep_interval)(self, *args, **kwargs)

    def __getattr__(self, name):
        func = ReplicaSetConnection.__getattr__(self, name)

        if isinstance(func, types.FunctionType):
            return wrap(func, self._num_retries, self._sleep_interval)
        else:
            return func

    def __getitem__(self, name):
        return self.__getattr__(name)

def main():
    # A simple regression test
    from pychartbeat.loggingutils import basicConfig
    from pychartbeat.options import define, options, parse_command_line

    define("console", default=True, type=bool)
    define("sleep_interval", default=1, type=int)
    define("num_retries", default=1, type=int)

    parse_command_line()
    basicConfig(options=options)

    for i in range(1000):
        c = AutoReconnectingConnection('localhost',
                num_retries=options.num_retries,
                sleep_interval=options.sleep_interval,
                )
        logging.info('save: %s' % c['test']['conn'].save({'idx':i}))
        logging.info('found: %s' % c['test']['conn'].find_one({'idx':i}))
        time.sleep(1)

if __name__ == '__main__':
    main()