nilmdb/nilmdb/client/httpclient.py

"""HTTP client library"""

import nilmdb
import nilmdb.utils
from nilmdb.client.errors import ClientError, ServerError, Error

import simplejson as json
import urlparse
import pycurl
import cStringIO

class HTTPClient(object):
    """Class to manage and perform HTTP requests from the client"""
    def __init__(self, baseurl = ""):
        """If baseurl is supplied, all other functions that take
        a URL can be given a relative URL instead."""
        # Verify / clean up URL
        reparsed = urlparse.urlparse(baseurl).geturl()
        if '://' not in reparsed:
            reparsed = urlparse.urlparse("http://" + baseurl).geturl()
        self.baseurl = reparsed
        self.curl = pycurl.Curl()
        self.curl.setopt(pycurl.SSL_VERIFYHOST, 2)
        self.curl.setopt(pycurl.FOLLOWLOCATION, 1)
        self.curl.setopt(pycurl.MAXREDIRS, 5)
        self._setup_url()

    def _setup_url(self, url = "", params = ""):
        url = urlparse.urljoin(self.baseurl, url)
        if params:
            url = urlparse.urljoin(
                url, "?" + nilmdb.utils.urllib.urlencode(params))
        self.curl.setopt(pycurl.URL, url)
        self.url = url

    def _check_error(self, body = None):
        code = self.curl.getinfo(pycurl.RESPONSE_CODE)
        if code == 200:
            return
        # Default variables for exception.  We use the entire body as
        # the default message, in case we can't extract it from a JSON
        # response.
        args = { "url" : self.url,
                 "status" : str(code),
                 "message" : body,
                 "traceback" : None }
        try:
            # Fill with server-provided data if we can
            jsonerror = json.loads(body)
            args["status"] = jsonerror["status"]
            args["message"] = jsonerror["message"]
            args["traceback"] = jsonerror["traceback"]
        except Exception: # pragma: no cover
            pass
        if code >= 400 and code <= 499:
            raise ClientError(**args)
        else: # pragma: no cover
            if code >= 500 and code <= 599:
                if args["message"] is None:
                    args["message"] = ("(no message; try disabling " +
                                       "response.stream option in " +
                                       "nilmdb.server for better debugging)")
                raise ServerError(**args)
            else:
                raise Error(**args)

    def _req_generator(self, url, params):
        """
        Like self._req(), but runs the perform in a separate thread.
        It returns a generator that spits out arbitrary-sized chunks
        of the resulting data, instead of using the WRITEFUNCTION
        callback.
        """
        self._setup_url(url, params)
        self._status = None
        error_body = ""
        self._headers = ""
        def header_callback(data):
            if self._status is None:
                self._status = int(data.split(" ")[1])
            self._headers += data
        self.curl.setopt(pycurl.HEADERFUNCTION, header_callback)
        def perform(callback):
            self.curl.setopt(pycurl.WRITEFUNCTION, callback)
            self.curl.perform()
        try:
            with nilmdb.utils.Iteratorizer(perform, curl_hack = True) as it:
                for i in it:
                    if self._status == 200:
                        # If we had a 200 response, yield the data to caller.
                        yield i
                    else:
                        # Otherwise, collect it into an error string.
                        error_body += i
        except pycurl.error as e:
            raise ServerError(status = "502 Error",
                              url = self.url,
                              message = e[1])
        # Raise an exception if there was an error
        self._check_error(error_body)

    def _req(self, url, params):
        """
        GET or POST that returns raw data.  Returns the body
        data as a string, or raises an error if it contained an error.
        """
        self._setup_url(url, params)
        body = cStringIO.StringIO()
        self.curl.setopt(pycurl.WRITEFUNCTION, body.write)
        self._headers = ""
        def header_callback(data):
            self._headers += data
        self.curl.setopt(pycurl.HEADERFUNCTION, header_callback)
        try:
            self.curl.perform()
        except pycurl.error as e:
            raise ServerError(status = "502 Error",
                              url = self.url,
                              message = e[1])
        body_str = body.getvalue()
        # Raise an exception if there was an error
        self._check_error(body_str)
        return body_str

    def close(self):
        self.curl.close()

    def _iterate_lines(self, it):
        """
        Given an iterator that returns arbitrarily-sized chunks
        of data, return '\n'-delimited lines of text
        """
        partial = ""
        for chunk in it:
            partial += chunk
            lines = partial.split("\n")
            for line in lines[0:-1]:
                yield line
            partial = lines[-1]
        if partial != "":
            yield partial

    # Non-generator versions
    def _doreq(self, url, params, retjson):
        """
        Perform a request, and return the body.

        url: URL to request (relative to baseurl)
        params: dictionary of query parameters
        retjson: expect JSON and return python objects instead of string
        """
        out = self._req(url, params)
        if retjson:
            return json.loads(out)
        return out

    def get(self, url, params = None, retjson = True):
        """Simple GET"""
        self.curl.setopt(pycurl.UPLOAD, 0)
        return self._doreq(url, params, retjson)

    def put(self, url, postdata, params = None, retjson = True):
        """Simple PUT"""
        self.curl.setopt(pycurl.UPLOAD, 1)
        self._setup_url(url, params)
        data = cStringIO.StringIO(postdata)
        self.curl.setopt(pycurl.READFUNCTION, data.read)
        return self._doreq(url, params, retjson)

    # Generator versions
    def _doreq_gen(self, url, params, retjson):
        """
        Perform a request, and return lines of the body in a generator.

        url: URL to request (relative to baseurl)
        params: dictionary of query parameters
        retjson: expect JSON and yield python objects instead of strings
        """
        for line in self._iterate_lines(self._req_generator(url, params)):
            if retjson:
                yield json.loads(line)
            else:
                yield line

    def get_gen(self, url, params = None, retjson = True):
        """Simple GET, returning a generator"""
        self.curl.setopt(pycurl.UPLOAD, 0)
        return self._doreq_gen(url, params, retjson)

    def put_gen(self, url, postdata, params = None, retjson = True):
        """Simple PUT, returning a generator"""
        self.curl.setopt(pycurl.UPLOAD, 1)
        self._setup_url(url, params)
        data = cStringIO.StringIO(postdata)
        self.curl.setopt(pycurl.READFUNCTION, data.read)
        return self._doreq_gen(url, params, retjson)