gistfile1.py

def retry_on_connection_error(retries=2, sleep=0.1):
    """ Retries on Autoreconnect errors. Wrap functions you want to ensure they
    retry if there's a connection failure, giving the chance to mongo to failover.

    params:
        retries: the number of retries
        sleep: seconds to sleep between retries

    >>> @retry_on_connection_error(retries=1)
    ... def super_sensitive_function(data=2)
    ...     return 2 ** 2

    """
    def wrapper(fn):
        @functools.wraps(fn)
        def wrapped_fn(*args, **kwargs):
            max_retries = retries
            import time
            while 1:
                try:
                    return fn(*args, **kwargs)
                except AutoReconnect:
                    if max_retries == 0:
                        raise
                    else:
                        max_retries -= 1
                        time.sleep(sleep)
        return wrapped_fn
    return wrapper