Support for HTTP auth. Thanks, nkvoll!

This commit is contained in:
Honza Kral
2013-09-24 22:55:04 +02:00
parent 8b7c7190f1
commit 7a440fda4a
2 changed files with 42 additions and 5 deletions
+23 -4
View File
@@ -10,10 +10,20 @@ from .base import Connection
from ..exceptions import ConnectionError from ..exceptions import ConnectionError
class RequestsHttpConnection(Connection): 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) super(RequestsHttpConnection, self).__init__(host=host, port=port, **kwargs)
self.session = requests.session() 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): def perform_request(self, method, url, params=None, body=None, timeout=None):
url = self.host + self.url_prefix + url url = self.host + self.url_prefix + url
@@ -41,10 +51,19 @@ class RequestsHttpConnection(Connection):
class Urllib3HttpConnection(Connection): class Urllib3HttpConnection(Connection):
""" """
Default connection class using the `urllib3` library and the http protocol. 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) 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): def perform_request(self, method, url, params=None, body=None, timeout=None):
url = self.url_prefix + url url = self.url_prefix + url
+19 -1
View File
@@ -2,10 +2,20 @@ import re
from mock import Mock, patch from mock import Mock, patch
from elasticsearch.exceptions import TransportError from elasticsearch.exceptions import TransportError
from elasticsearch.connection import RequestsHttpConnection from elasticsearch.connection import RequestsHttpConnection, Urllib3HttpConnection
from .test_cases import TestCase 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): class TestRequestsConnection(TestCase):
def _get_mock_connection(self, connection_params={}, status_code=200, response_body=u'{}'): def _get_mock_connection(self, connection_params={}, status_code=200, response_body=u'{}'):
con = RequestsHttpConnection(**connection_params) con = RequestsHttpConnection(**connection_params)
@@ -31,6 +41,14 @@ class TestRequestsConnection(TestCase):
self.assertEquals(1, len(args)) self.assertEquals(1, len(args))
return args[0] 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): def test_repr(self):
con = self._get_mock_connection({"host": "elasticsearch.com", "port": 443}) con = self._get_mock_connection({"host": "elasticsearch.com", "port": 443})
self.assertEquals('<RequestsHttpConnection: http://elasticsearch.com:443>', repr(con)) self.assertEquals('<RequestsHttpConnection: http://elasticsearch.com:443>', repr(con))