python-consul · luhhujbb · Sep 23, 2019 · Sep 23, 2019 · Sep 23, 2019 · Sep 24, 2019
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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):