Added Transport class that manages connections

elasticsearch/transport.py

@@ -0,0 +1,58 @@
+from .connection import RequestsHttpConnection
+from .connection_pool import ConnectionPool
+from .serializer import JSONSerializer
+from .exceptions import TransportError
+
+class Transport(object):
+    def __init__(self, hosts, connection_class=RequestsHttpConnection,
+        connection_pool_class=ConnectionPool, serializer=JSONSerializer(),
+        max_retries=3, **kwargs):
+
+        self.max_retries = 3
+
+        # data serializer
+        self.serializer = serializer
+
+        # store all strategies...
+        self.connection_pool_class = connection_pool_class
+        self.connection_class = connection_class
+
+        # ...save kwargs to be passed to the connections
+        self.kwargs = kwargs
+        self.hosts = hosts
+
+        # ...and instantiate them
+        self.set_connections(hosts)
+
+    def add_connection(self, host):
+        self.hosts.append(host)
+        self.set_connections(self.hosts)
+
+    def set_connections(self, hosts):
+        # construct the connections
+        def _create_connection(host):
+            kwargs = self.kwargs.copy()
+            kwargs.update(host)
+            return self.connection_class(**kwargs)
+        connections = list(map(_create_connection, hosts))
+
+        # pass the hosts dicts to the connection pool to optionally extract parameters from
+        self.connection_pool = self.connection_pool_class(zip(connections, hosts), **self.kwargs)
+
+    def perform_request(self, method, url, params=None, body=None):
+        for attempt in range(self.max_retries):
+            connection = self.connection_pool.get_connection()
+
+            if body:
+                body = self.serializer.dumps(body)
+            try:
+                status, raw_data = connection.perform_request(method, url, params, body)
+            except TransportError:
+                self.connection_pool.mark_dead(connection)
+
+                # raise exception on last retry
+                if attempt + 1 == self.max_retries:
+                    raise
+            else:
+                return status, self.serializer.loads(raw_data)
+

test_elasticsearch/test_transport.py

@@ -0,0 +1,57 @@
+from unittest import TestCase
+
+from elasticsearch.transport import Transport
+from elasticsearch.connection import Connection
+from elasticsearch.exceptions import TransportError
+
+class DummyConnection(Connection):
+    def __init__(self, **kwargs):
+        self.exception = kwargs.pop('exception', None)
+        self.status, self.data = kwargs.pop('status', 200), kwargs.pop('data', '{}')
+        self.calls = []
+        super(DummyConnection, self).__init__(**kwargs)
+
+    def perform_request(self, *args, **kwargs):
+        self.calls.append((args, kwargs))
+        if self.exception:
+            raise self.exception
+        return self.status, self.data
+
+class TestTransport(TestCase):
+    def test_kwargs_passed_on_to_connections(self):
+        t = Transport([{'host': 'google.com'}], port=123)
+        self.assertEquals(1, len(t.connection_pool.connections))
+        self.assertEquals('http://google.com:123', t.connection_pool.connections[0].host)
+
+    def test_kwargs_passed_on_to_connection_pool(self):
+        dt = object()
+        t = Transport([{}], dead_timeout=dt)
+        self.assertIs(dt, t.connection_pool.dead_timeout)
+
+    def test_custom_connection_class(self):
+        class MyConnection(object):
+            def __init__(self, **kwargs):
+                self.kwargs = kwargs
+        t = Transport([{}], connection_class=MyConnection)
+        self.assertEquals(1, len(t.connection_pool.connections))
+        self.assertIsInstance(t.connection_pool.connections[0], MyConnection)
+
+    def test_add_connection(self):
+        t = Transport([{}], randomize_hosts=False)
+        t.add_connection({"host": "google.com"})
+
+        self.assertEquals(2, len(t.connection_pool.connections))
+        self.assertEquals('http://google.com:9200', t.connection_pool.connections[1].host)
+
+    def test_request_will_fail_after_X_retries(self):
+        t = Transport([{'exception': TransportError('abandon ship')}], connection_class=DummyConnection)
+
+        self.assertRaises(TransportError, t.perform_request, 'GET', '/')
+        self.assertEquals(3, len(t.connection_pool.get_connection().calls))
+
+    def test_failed_connection_will_be_marked_as_dead(self):
+        t = Transport([{'exception': TransportError('abandon ship')}], connection_class=DummyConnection)
+
+        self.assertRaises(TransportError, t.perform_request, 'GET', '/')
+        self.assertEquals(0, len(t.connection_pool.connections))
+