Add streaming get

Author: gmr

consulate/adapters.py

@@ -4,15 +4,16 @@
 """
 import json
 import logging
+import socket
 
 import requests
+import requests.exceptions
 try:
     import requests_unixsocket
 except ImportError:  # pragma: no cover
     requests_unixsocket = None
 
-from consulate import api
-from consulate import utils
+from consulate import api, exceptions, utils
 
 LOGGER = logging.getLogger(__name__)
 
@@ -71,23 +72,51 @@ def delete(self, uri):
         return self._process_response(
             self.session.delete(uri, timeout=self.timeout))
 
-    def get(self, uri):
+    def get(self, uri, timeout=None):
         """Perform a HTTP get
 
         :param src uri: The URL to send the DELETE to
+        :param timeout: How long to wait on the response
+        :type timeout: int or float or None
         :rtype: consulate.api.Response
 
         """
         LOGGER.debug("GET %s", uri)
-        return self._process_response(
-            self.session.get(uri, timeout=self.timeout))
+        try:
+            return self._process_response(
+                self.session.get(uri, timeout=timeout or self.timeout))
+        except (requests.exceptions.ConnectionError,
+                OSError, socket.error) as err:
+            raise exceptions.RequestError(str(err))
+
+    def get_stream(self, uri):
+        """Perform a HTTP get that returns the response as a stream.
+
+        :param src uri: The URL to send the DELETE to
+        :rtype: iterator
+
+        """
+        LOGGER.debug("GET Stream from %s", uri)
+        try:
+            response = self.session.get(uri, stream=True)
+        except (requests.exceptions.ConnectionError,
+                OSError, socket.error) as err:
+            raise exceptions.RequestError(str(err))
+        if response.encoding is None:
+            response.encoding = 'utf-8'
+        if utils.response_ok(response):
+            for line in response.iter_lines():
+                if line:
+                    yield line.decode('utf-8')
 
     @prepare_data
-    def put(self, uri, data=None):
+    def put(self, uri, data=None, timeout=None):
         """Perform a HTTP put
 
         :param src uri: The URL to send the DELETE to
         :param str data: The PUT data
+        :param timeout: How long to wait on the response
+        :type timeout: int or float or None
         :rtype: consulate.api.Response
 
         """
@@ -96,9 +125,14 @@ def put(self, uri, data=None):
             'Content-Type': CONTENT_FORM
             if utils.is_string(data) else CONTENT_JSON
         }
-        return self._process_response(
-            self.session.put(
-                uri, data=data, headers=headers, timeout=self.timeout))
+        try:
+            return self._process_response(
+                self.session.put(
+                    uri, data=data, headers=headers,
+                    timeout=timeout or self.timeout))
+        except (requests.exceptions.ConnectionError,
+                OSError, socket.error) as err:
+            raise exceptions.RequestError(str(err))
 
     @staticmethod
     def _process_response(response):
@@ -109,8 +143,12 @@ def _process_response(response):
         :rtype: consulate.api.Response
 
         """
-        return api.Response(
-            response.status_code, response.content, response.headers)
+        try:
+            return api.Response(
+                response.status_code, response.content, response.headers)
+        except (requests.exceptions.ConnectionError,
+                OSError, socket.error) as err:
+            raise exceptions.RequestError(str(err))
 
 
 class UnixSocketRequest(Request):

consulate/api/base.py

@@ -51,22 +51,22 @@ def _build_uri(self, params, query_params=None):
                                         urlencode(query_params))
         return '{0}/{1}'.format(self._base_uri, path)
 
-    def _get(self, params, query_params=None, raise_on_404=False):
+    def _get(self, params, query_params=None, raise_on_404=False,
+             timeout=None):
         """Perform a GET request
 
         :param list params: List of path parts
         :param dict query_params: Build query parameters
+        :param timeout: How long to wait on the request for
+        :type timeout: int or float or None
 
         """
-        response = self._adapter.get(self._build_uri(params, query_params))
-        if response.status_code == 200:
+        response = self._adapter.get(
+            self._build_uri(params, query_params), timeout=timeout)
+        if not response:
+            print(response)
+        if utils.response_ok(response, raise_on_404):
             return response.body
-        elif response.status_code == 401:
-            raise exceptions.ACLDisabled(response.body)
-        elif response.status_code == 403:
-            raise exceptions.Forbidden(response.body)
-        elif response.status_code == 404 and raise_on_404:
-            raise exceptions.NotFound(response.body)
         return []
 
     def _get_list(self, params, query_params=None):
@@ -81,20 +81,36 @@ def _get_list(self, params, query_params=None):
             return [result]
         return result
 
+    def _get_stream(self, params, query_params=None):
+        """Return a list queried from Consul
+
+        :param list params: List of path parts
+        :param dict query_params: Build query parameters
+        :rtype: iterator
+
+        """
+        for line in self._adapter.get_stream(
+                self._build_uri(params, query_params)):
+            yield line
+
     def _get_no_response_body(self, url_parts, query=None):
-        return self._adapter.get(self._build_uri(url_parts,
-                                                 query)).status_code == 200
+        return utils.response_ok(
+            self._adapter.get(self._build_uri(url_parts, query)))
 
     def _get_response_body(self, url_parts, query=None):
-        return self._adapter.get(self._build_uri(url_parts, query)).body
+        response = self._adapter.get(self._build_uri(url_parts, query))
+        if utils.response_ok(response):
+            return response.body
 
     def _put_no_response_body(self, url_parts, query=None, payload=None):
-        return self._adapter.put(self._build_uri(url_parts, query),
-                                 payload).status_code == 200
+        return utils.response_ok(
+            self._adapter.put(self._build_uri(url_parts, query), payload))
 
     def _put_response_body(self, url_parts, query=None, payload=None):
-        return self._adapter.put(self._build_uri(url_parts, query),
-                                 payload).body
+        response = self._adapter.put(
+            self._build_uri(url_parts, query), payload)
+        if utils.response_ok(response):
+            return response.body
 
 
 class Response(object):