diff --git a/consul/aio.py b/consul/aio.py index 5e347f53..7274e7cf 100644 --- a/consul/aio.py +++ b/consul/aio.py @@ -25,9 +25,10 @@ def __init__(self, *args, loop=None, **kwargs): def _request(self, callback, method, uri, data=None): resp = yield from self._session.request(method, uri, data=data) body = yield from resp.text(encoding='utf-8') + content = yield from resp.read() if resp.status == 599: raise base.Timeout - r = base.Response(resp.status, resp.headers, body) + r = base.Response(resp.status, resp.headers, body, content) return callback(r) # python prior 3.4.1 does not play nice with __del__ method diff --git a/consul/base.py b/consul/base.py index ee6ab254..885bca8a 100755 --- a/consul/base.py +++ b/consul/base.py @@ -162,7 +162,8 @@ def _compat( return ret -Response = collections.namedtuple('Response', ['code', 'headers', 'body']) +Response = collections.namedtuple('Response', + ['code', 'headers', 'body', 'content']) # @@ -244,6 +245,16 @@ def cb(response): return data return cb + @classmethod + def binary(klass): + """ + This method simply returns response body, usefull for snapshot + """ + def cb(response): + CB._status(response) + return response.content + return cb + class HTTPClient(six.with_metaclass(abc.ABCMeta, object)): def __init__(self, host='127.0.0.1', port=8500, scheme='http', @@ -343,6 +354,7 @@ def __init__( self.query = Consul.Query(self) self.coordinate = Consul.Coordinate(self) self.operator = Consul.Operator(self) + self.snapshot = Consul.Snapshot(self) class Event(object): """ @@ -2434,3 +2446,50 @@ def raft_config(self): """ return self.agent.http.get( CB.json(), '/v1/operator/raft/configuration') + + class Snapshot(object): + def __init__(self, agent): + self.agent = agent + + def get(self, dc=None, token=None): + """ + Returns gzipped snapshot of current consul cluster + """ + params = [] + token = token or self.agent.token + if token: + params.append(('token', token)) + if dc: + params.append(('dc', dc)) + return self.agent.http.get( + CB.binary(), '/v1/snapshot', params=params) + + def save(self, file_path, dc=None, token=None): + """ + Backup snapshot in a file + """ + backup_file = open(file_path, 'w+b') + backup_file.write(self.get(dc, token)) + backup_file.close() + + def restore(self, file_path=None, data=None, dc=None, token=None): + """ + Restore snapshot from a file or from data + """ + if file_path or data: + if file_path: + backup_file = open(file_path, 'rb') + data = backup_file.read() + backup_file.close() + params = [] + token = token or self.agent.token + if token: + params.append(('token', token)) + if dc: + params.append(('dc', dc)) + res = self.agent.http.put(CB.bool(), '/v1/snapshot', + params=params, + data=data) + return res + else: + return False diff --git a/consul/std.py b/consul/std.py index 96a5b9dc..f2a77a61 100644 --- a/consul/std.py +++ b/consul/std.py @@ -14,7 +14,8 @@ def __init__(self, *args, **kwargs): def response(self, response): response.encoding = 'utf-8' return base.Response( - response.status_code, response.headers, response.text) + response.status_code, response.headers, + response.text, response.content) def get(self, callback, path, params=None): uri = self.uri(path, params) diff --git a/consul/tornado.py b/consul/tornado.py index 50507748..3ae4dd8e 100644 --- a/consul/tornado.py +++ b/consul/tornado.py @@ -16,7 +16,8 @@ def __init__(self, *args, **kwargs): def response(self, response): return base.Response( - response.code, response.headers, response.body.decode('utf-8')) + response.code, response.headers, + response.body.decode('utf-8'), response.body) @gen.coroutine def _request(self, callback, request): diff --git a/consul/twisted.py b/consul/twisted.py index 206105cb..3a7a20e4 100644 --- a/consul/twisted.py +++ b/consul/twisted.py @@ -47,8 +47,8 @@ def __init__(self, contextFactory, *args, **kwargs): self.client = TreqHTTPClient(Agent(**agent_kwargs)) @staticmethod - def response(code, headers, text): - return base.Response(code, headers, text) + def response(code, headers, text, content): + return base.Response(code, headers, text, content) @staticmethod def compat_string(value): @@ -70,7 +70,8 @@ def _get_resp(self, response): for k, v in dict(response.headers.getAllRawHeaders()).items() ]) body = yield response.text(encoding='utf-8') - returnValue((response.code, headers, body)) + content = yield response.content() + returnValue((response.code, headers, body, content)) @inlineCallbacks def request(self, callback, method, url, **kwargs): diff --git a/tests/test_base.py b/tests/test_base.py index ef76a9fa..1183cbe4 100755 --- a/tests/test_base.py +++ b/tests/test_base.py @@ -131,39 +131,39 @@ def test_meta(self): class TestCB(object): def test_status_200_passes(self): - response = consul.base.Response(200, None, None) + response = consul.base.Response(200, None, None, None) CB._status(response) @pytest.mark.parametrize( 'response, expected_exception', [ - (Response(400, None, None), consul.base.BadRequest), - (Response(401, None, None), consul.base.ACLDisabled), - (Response(403, None, None), consul.base.ACLPermissionDenied), + (Response(400, None, None, None), consul.base.BadRequest), + (Response(401, None, None, None), consul.base.ACLDisabled), + (Response(403, None, None, None), consul.base.ACLPermissionDenied), ]) def test_status_4xx_raises_error(self, response, expected_exception): with pytest.raises(expected_exception): CB._status(response) def test_status_404_allow_404(self): - response = Response(404, None, None) + response = Response(404, None, None, None) CB._status(response, allow_404=True) def test_status_404_dont_allow_404(self): - response = Response(404, None, None) + response = Response(404, None, None, None) with pytest.raises(consul.base.NotFound): CB._status(response, allow_404=False) def test_status_405_raises_generic_ClientError(self): - response = Response(405, None, None) + response = Response(405, None, None, None) with pytest.raises(consul.base.ClientError): CB._status(response) @pytest.mark.parametrize( 'response', [ - Response(500, None, None), - Response(599, None, None), + Response(500, None, None, None), + Response(599, None, None, None), ]) def test_status_5xx_raises_error(self, response): with pytest.raises(consul.base.ConsulException):