tornadoes wrapper for multiple connections

tornadoes_multiconn.py

"""
This tornadoes wrappers wraps around ESConnection from tornadoes

This allow us to provide multiple host
"""
import tornado.gen

import tornadoes
import re
import random

HOST_REGEX = re.compile("(?P<host>[^:]+):(?P<port>\d+)")

class ESMultiConnection(object):

    def __init__(self, hosts, max_retries=3, **kwargs):
        """
        hosts               each host should be in the format <host>:<port>

        the rest of the arguments is the same as the original ESConnection

        host and port will be popped from kwargs
        """
        kwargs.pop("host", None)
        kwargs.pop("port", None)
        self.es_connections = []
        for host in hosts:
            match = HOST_REGEX.match(host)
            if not match:
                raise Exception("Invalid host {0}".format(host))
            self.es_connections.append(tornadoes.ESConnection(**match.groupdict(), **kwargs))

        self.max_retries = max_retries
        for func in [ "search", "apply_search", "get_by_path", "get", "put", "update", "delete",
                "count" ]:
            setattr(self, func, self._get_wrapped_function(func))

    def _get_wrapped_function(self, func_name):

        @tornado.gen.coroutine
        def _func(**kwargs):
            es_conns = [ i for i in self.es_connections ]
            random.shuffle(es_conns)
            response = None

            for conn in es_conns[:self.max_retries]:
                func = getattr(conn, func_name)
                response = yield func(**kwargs)
                # any exception from the func will be raise out, which is fine because any exception
                # will most likely affect other es_conn as well.
                if response.code == 599:
                    print("Connection fail for {0}".format(conn.url))
                    continue

                return response

            return response

        return _func