Created
May 23, 2012 19:47
-
-
Save daniyalzade/2777345 to your computer and use it in GitHub Desktop.
AutoReconnectingConnection
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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() |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment