diff --git a/pydruid/client.py b/pydruid/client.py index 11fe40a7..d0ef02e2 100755 --- a/pydruid/client.py +++ b/pydruid/client.py @@ -22,13 +22,19 @@ from six.moves import urllib from pydruid.query import QueryBuilder - +from base64 import b64encode class BaseDruidClient(object): def __init__(self, url, endpoint): self.url = url self.endpoint = endpoint self.query_builder = QueryBuilder() + self.username = None + self.password = None + + def set_basic_auth_credentials(self, username, password): + self.username = username + self.password = password def _prepare_url_headers_and_body(self, query): querystr = json.dumps(query.query_dict).encode('utf-8') @@ -37,6 +43,11 @@ def _prepare_url_headers_and_body(self, query): else: url = self.url + '/' + self.endpoint headers = {'Content-Type': 'application/json'} + if (self.username is not None) and (self.password is not None): + username_password = \ + b64encode(bytes('{}:{}'.format(self.username, self.password))) + headers['Authorization'] = 'Basic {}'.format(username_password) + return headers, querystr, url def _post(self, query): @@ -481,8 +492,6 @@ def _post(self, query): err = json.loads(e.read().decode("utf-8")) except (ValueError, AttributeError, KeyError): pass - else: - err = err.get('error', None) raise IOError('{0} \n Druid Error: {1} \n Query is: {2}'.format( e, err, json.dumps(query.query_dict, indent=4)))