diff --git a/elasticsearch/connection/http.py b/elasticsearch/connection/http.py index 9e15e43e..d6399848 100644 --- a/elasticsearch/connection/http.py +++ b/elasticsearch/connection/http.py @@ -15,18 +15,23 @@ class RequestsHttpConnection(Connection): :arg http_auth: optional http auth information as either ':' separated string or a tuple + :arg use_https: use https for the connection if `True` """ - def __init__(self, host='localhost', port=9200, http_auth=None, **kwargs): + def __init__(self, host='localhost', port=9200, http_auth=None, use_https=False, **kwargs): super(RequestsHttpConnection, self).__init__(host=host, port=port, **kwargs) self.session = requests.session() - if http_auth is not None: + if http_auth is not None: if not isinstance(http_auth, tuple): http_auth = tuple(http_auth.split(':', 1)) self.session.auth = http_auth + self.base_url = 'http%s://%s:%d%s' % ( + 's' if use_https else '', + host, port, self.url_prefix + ) def perform_request(self, method, url, params=None, body=None, timeout=None): - url = self.host + self.url_prefix + url + url = self.base_url + url # use prepared requests so that requests formats url and params for us to log request = requests.Request(method, url, params=params or {}, data=body).prepare() @@ -54,8 +59,9 @@ class Urllib3HttpConnection(Connection): :arg http_auth: optional http auth information as either ':' separated string or a tuple + :arg use_https: use https for the connection if `True` """ - def __init__(self, host='localhost', port=9200, http_auth=None, **kwargs): + def __init__(self, host='localhost', port=9200, http_auth=None, use_https=False, **kwargs): super(Urllib3HttpConnection, self).__init__(host=host, port=port, **kwargs) headers = {} if http_auth is not None: @@ -63,7 +69,11 @@ class Urllib3HttpConnection(Connection): http_auth = ':'.join(http_auth) headers = urllib3.make_headers(basic_auth=http_auth) - self.pool = urllib3.HTTPConnectionPool(host, port=port, timeout=kwargs.get('timeout', None), headers=headers) + pool_class = urllib3.HTTPConnectionPool + if use_https: + pool_class = urllib3.HTTPSConnectionPool + + self.pool = pool_class(host, port=port, timeout=kwargs.get('timeout', None), headers=headers) def perform_request(self, method, url, params=None, body=None, timeout=None): url = self.url_prefix + url diff --git a/test_elasticsearch/test_connection.py b/test_elasticsearch/test_connection.py index c1f68719..f5cefec7 100644 --- a/test_elasticsearch/test_connection.py +++ b/test_elasticsearch/test_connection.py @@ -1,5 +1,6 @@ import re from mock import Mock, patch +import urllib3 from elasticsearch.exceptions import TransportError from elasticsearch.connection import RequestsHttpConnection, Urllib3HttpConnection @@ -15,6 +16,13 @@ class TestUrllib3Connection(TestCase): con = Urllib3HttpConnection(http_auth=('username', 'secret')) self.assertEquals({'authorization': 'Basic dXNlcm5hbWU6c2VjcmV0'}, con.pool.headers) + def test_uses_https_if_specified(self): + con = Urllib3HttpConnection(use_https=True) + self.assertIsInstance(con.pool, urllib3.HTTPSConnectionPool) + + def test_doesnt_use_https_if_not_specified(self): + con = Urllib3HttpConnection() + self.assertIsInstance(con.pool, urllib3.HTTPConnectionPool) class TestRequestsConnection(TestCase): def _get_mock_connection(self, connection_params={}, status_code=200, response_body=u'{}'): @@ -41,6 +49,14 @@ class TestRequestsConnection(TestCase): self.assertEquals(1, len(args)) return args[0] + def test_use_https_if_specified(self): + con = self._get_mock_connection({'use_https': True, 'url_prefix': 'url'}) + request = self._get_request(con, 'GET', '/') + + self.assertEquals('https://localhost:9200/url/', request.url) + self.assertEquals('GET', request.method) + self.assertEquals(None, request.body) + def test_http_auth(self): con = RequestsHttpConnection(http_auth='username:secret') self.assertEquals(('username', 'secret'), con.session.auth)