diff --git a/patroni/dcs/consul.py b/patroni/dcs/consul.py index ba6799af..a9acb08b 100644 --- a/patroni/dcs/consul.py +++ b/patroni/dcs/consul.py @@ -8,6 +8,7 @@ import ssl import time import urllib3 +from collections import namedtuple from consul import ConsulException, NotFound, base from urllib3.exceptions import HTTPError from six.moves.urllib.parse import urlencode, urlparse, quote @@ -36,6 +37,9 @@ class InvalidSession(ConsulException): """invalid session""" +Response = namedtuple('Response', 'code,headers,body,content') + + class HTTPClient(object): def __init__(self, host='127.0.0.1', port=8500, token=None, scheme='http', verify=True, cert=None, ca_cert=None): @@ -71,16 +75,17 @@ class HTTPClient(object): @staticmethod def response(response): - data = response.data.decode('utf-8') + content = response.data + body = content.decode('utf-8') if response.status == 500: - msg = '{0} {1}'.format(response.status, data) - if data.startswith('Invalid Session TTL'): + msg = '{0} {1}'.format(response.status, body) + if body.startswith('Invalid Session TTL'): raise InvalidSessionTTL(msg) - elif data.startswith('invalid session'): + elif body.startswith('invalid session'): raise InvalidSession(msg) else: raise ConsulInternalError(msg) - return base.Response(response.status, response.headers, data) + return Response(response.status, response.headers, body, content) def uri(self, path, params=None): return '{0}{1}{2}'.format(self.base_uri, path, params and '?' + urlencode(params) or '') @@ -89,7 +94,7 @@ class HTTPClient(object): if method not in ('get', 'post', 'put', 'delete'): raise AttributeError("HTTPClient instance has no attribute '{0}'".format(method)) - def wrapper(callback, path, params=None, data=''): + def wrapper(callback, path, params=None, data='', headers=None): # python-consul doesn't allow to specify ttl smaller then 10 seconds # because session_ttl_min defaults to 10s, so we have to do this ugly dirty hack... if method == 'put' and path == '/v1/session/create': @@ -110,8 +115,9 @@ class HTTPClient(object): kwargs['timeout'] = timeout + max(timeout/15.0, 1) else: kwargs['timeout'] = self._read_timeout + kwargs['headers'] = (headers or {}).copy() + kwargs['headers'].update(urllib3.make_headers(user_agent=USER_AGENT)) token = params.pop('token', self.token) if isinstance(params, dict) else self.token - kwargs['headers'] = urllib3.make_headers(user_agent=USER_AGENT) if token: kwargs['headers']['X-Consul-Token'] = token return callback(self.response(self.http.request(method.upper(), self.uri(path, params), **kwargs))) @@ -126,7 +132,7 @@ class ConsulClient(base.Consul): self.token = kwargs.get('token') super(ConsulClient, self).__init__(*args, **kwargs) - def connect(self, *args, **kwargs): + def http_connect(self, *args, **kwargs): kwargs.update(dict(zip(['host', 'port', 'scheme', 'verify'], args))) if self._cert: kwargs['cert'] = self._cert @@ -136,6 +142,9 @@ class ConsulClient(base.Consul): kwargs['token'] = self.token return HTTPClient(**kwargs) + def connect(self, *args, **kwargs): + return self.http_connect(*args, **kwargs) + def reload_config(self, config): self.http.token = self.token = config.get('token') self.consistency = config.get('consistency', 'default') diff --git a/tests/test_consul.py b/tests/test_consul.py index 6d0c4732..a38a2869 100644 --- a/tests/test_consul.py +++ b/tests/test_consul.py @@ -4,7 +4,7 @@ import unittest from consul import ConsulException, NotFound from mock import Mock, patch from patroni.dcs.consul import AbstractDCS, Cluster, Consul, ConsulInternalError, \ - ConsulError, HTTPClient, InvalidSessionTTL, InvalidSession + ConsulError, ConsulClient, HTTPClient, InvalidSessionTTL, InvalidSession from . import SleepException @@ -41,7 +41,8 @@ def kv_get(self, key, **kwargs): class TestHTTPClient(unittest.TestCase): def setUp(self): - self.client = HTTPClient('127.0.0.1', '8500', 'http', False) + c = ConsulClient() + self.client = c.http self.client.http.request = Mock() def test_get(self):