diff --git a/elasticsearch/connection/http.py b/elasticsearch/connection/http.py index c55be9f2..9e15e43e 100644 --- a/elasticsearch/connection/http.py +++ b/elasticsearch/connection/http.py @@ -10,10 +10,20 @@ from .base import Connection from ..exceptions import ConnectionError class RequestsHttpConnection(Connection): - """ Connection using the `requests` library. """ - def __init__(self, host='localhost', port=9200, **kwargs): + """ + Connection using the `requests` library. + + :arg http_auth: optional http auth information as either ':' separated + string or a tuple + """ + def __init__(self, host='localhost', port=9200, http_auth=None, **kwargs): super(RequestsHttpConnection, self).__init__(host=host, port=port, **kwargs) self.session = requests.session() + if http_auth is not None: + if not isinstance(http_auth, tuple): + http_auth = tuple(http_auth.split(':', 1)) + self.session.auth = http_auth + def perform_request(self, method, url, params=None, body=None, timeout=None): url = self.host + self.url_prefix + url @@ -41,10 +51,19 @@ class RequestsHttpConnection(Connection): class Urllib3HttpConnection(Connection): """ Default connection class using the `urllib3` library and the http protocol. + + :arg http_auth: optional http auth information as either ':' separated + string or a tuple """ - def __init__(self, host='localhost', port=9200, **kwargs): + def __init__(self, host='localhost', port=9200, http_auth=None, **kwargs): super(Urllib3HttpConnection, self).__init__(host=host, port=port, **kwargs) - self.pool = urllib3.HTTPConnectionPool(host, port=port, timeout=kwargs.get('timeout', None)) + headers = {} + if http_auth is not None: + if isinstance(http_auth, (tuple, list)): + 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) 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 f471a54f..c1f68719 100644 --- a/test_elasticsearch/test_connection.py +++ b/test_elasticsearch/test_connection.py @@ -2,10 +2,20 @@ import re from mock import Mock, patch from elasticsearch.exceptions import TransportError -from elasticsearch.connection import RequestsHttpConnection +from elasticsearch.connection import RequestsHttpConnection, Urllib3HttpConnection from .test_cases import TestCase +class TestUrllib3Connection(TestCase): + def test_http_auth(self): + con = Urllib3HttpConnection(http_auth='username:secret') + self.assertEquals({'authorization': 'Basic dXNlcm5hbWU6c2VjcmV0'}, con.pool.headers) + + def test_http_auth_tuple(self): + con = Urllib3HttpConnection(http_auth=('username', 'secret')) + self.assertEquals({'authorization': 'Basic dXNlcm5hbWU6c2VjcmV0'}, con.pool.headers) + + class TestRequestsConnection(TestCase): def _get_mock_connection(self, connection_params={}, status_code=200, response_body=u'{}'): con = RequestsHttpConnection(**connection_params) @@ -31,6 +41,14 @@ class TestRequestsConnection(TestCase): self.assertEquals(1, len(args)) return args[0] + def test_http_auth(self): + con = RequestsHttpConnection(http_auth='username:secret') + self.assertEquals(('username', 'secret'), con.session.auth) + + def test_http_auth_tuple(self): + con = RequestsHttpConnection(http_auth=('username', 'secret')) + self.assertEquals(('username', 'secret'), con.session.auth) + def test_repr(self): con = self._get_mock_connection({"host": "elasticsearch.com", "port": 443}) self.assertEquals('', repr(con))