nilmdb/nilmdb/httpclient.py

"""HTTP client library"""

from __future__ import absolute_import
from nilmdb.printf import *

import time
import sys
import re
import os
import simplejson as json
import urlparse
import urllib
import pycurl
import cStringIO

import nilmdb.iteratorizer

class Error(Exception):
    """Base exception for both ClientError and ServerError responses"""
    def __init__(self,
                 status = "Unspecified error",
                 message = None,
                 url = None,
                 traceback = None):
        Exception.__init__(self, status)
        self.status = status     # e.g. "400 Bad Request"
        self.message = message   # textual message from the server
        self.url = url           # URL we were requesting
        self.traceback = traceback # server traceback, if available
    def __str__(self):
        s = sprintf("[%s]", self.status)
        if self.message:
            s += sprintf(" %s", self.message)
        if self.url:
            s += sprintf(" (%s)", self.url)
        if self.traceback: # pragma: no cover
            s += sprintf("\nServer traceback:\n%s", self.traceback)
        return s
class ClientError(Error):
    pass
class ServerError(Error):
    pass

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, "?" + urllib.urlencode(params, True))
        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
        args = { "url" : self.url,
                 "status" : str(code),
                 "message" : None,
                 "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:
                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 = ""
        def header_callback(data):
            if self._status is None:
                self._status = int(data.split(" ")[1])
        self.curl.setopt(pycurl.HEADERFUNCTION, header_callback)
        def func(callback):
            self.curl.setopt(pycurl.WRITEFUNCTION, callback)
            self.curl.perform()
        try:
            for i in nilmdb.iteratorizer.Iteratorizer(func):
                if self._status == 200:
                    # If we had a 200 response, yield the data to the 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)
        def header_callback(data):
            pass
        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._setup_url(url, params)
        data = cStringIO.StringIO(postdata)
        self.curl.setopt(pycurl.UPLOAD, 1)
        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._setup_url(url, params)
        data = cStringIO.StringIO(postdata)
        self.curl.setopt(pycurl.UPLOAD, 1)
        self.curl.setopt(pycurl.READFUNCTION, data.read)
        return self._doreq_gen(url, params, retjson)