diff --git a/elasticsearch/transport.py b/elasticsearch/transport.py new file mode 100644 index 00000000..b3490a1e --- /dev/null +++ b/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) + diff --git a/test_elasticsearch/test_transport.py b/test_elasticsearch/test_transport.py new file mode 100644 index 00000000..3ed3778f --- /dev/null +++ b/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)) +