nilmdb/nilmdb/client.py

"""Client utilities for accessing NILM database via HTTP"""

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

version = "1.0"

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 MyCurl(object):
    """Curl wrapper for HTTP client requests"""
    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(self, url, params):
        """GET or POST that returns raw data"""
        self._setup_url(url, params)
        body = cStringIO.StringIO()
        self.curl.setopt(pycurl.WRITEFUNCTION, body.write)
        try:
            self.curl.perform()
        except pycurl.error as e:
            raise ServerError(status = "502 Error",
                              url = self.url,
                              message = e[1])
        body_str = body.getvalue()
        self._check_error(body_str)
        return body_str

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

    def getjson(self, url, params = None):
        """Simple GET that returns JSON string"""
        return json.loads(self.get(url, params))

    def putjson(self, url, postdata, params = None):
        """Simple GET that returns JSON string"""
        return json.loads(self.put(url, postdata, params))

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

    def put(self, url, postdata, params = None):
        """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._req(url, params)

class Client(object):
    """Main client interface to the Nilm database."""

    client_version = version

    def __init__(self, url):
        self.curl = MyCurl(url)

    def _json_param(self, data):
        """Return compact json-encoded version of parameter"""
        return json.dumps(data, separators=(',',':'))

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

    def geturl(self):
        """Return the URL we're using"""
        return self.curl.baseurl

    def version(self):
        """Return server version"""
        return self.curl.getjson("version")

    def dbpath(self):
        """Return server database path"""
        return self.curl.getjson("dbpath")

    def stream_list(self, path = None, layout = None):
        params = {}
        if path is not None:
            params["path"] = path
        if layout is not None:
            params["layout"] = layout
        return self.curl.getjson("stream/list", params)

    def stream_get_metadata(self, path, keys = None):
        params = { "path": path }
        if keys is not None:
            params["key"] = keys
        return self.curl.getjson("stream/get_metadata", params)

    def stream_set_metadata(self, path, data):
        """Set stream metadata from a dictionary, replacing all existing
        metadata."""
        params = {
            "path": path,
            "data": self._json_param(data)
            }
        return self.curl.getjson("stream/set_metadata", params)

    def stream_update_metadata(self, path, data):
        """Update stream metadata from a dictionary"""
        params = {
            "path": path,
            "data": self._json_param(data)
            }
        return self.curl.getjson("stream/update_metadata", params)

    def stream_create(self, path, layout):
        """Create a new stream"""
        params = { "path": path,
                   "layout" : layout }
        return self.curl.getjson("stream/create", params)

    def stream_insert(self, path, data):
        """Insert data into a stream.  data should be a file-like object
        that provides ASCII data that matches the database layout for path."""
        params = { "path": path }

        # See design.md for a discussion of how much data to send.
        # These are soft limits -- actual data might be rounded up.
        max_data = 1048576
        max_time = 30

        def sendit():
            return self.curl.putjson("stream/insert", send_data, params)

        result = None
        start = time.time()
        send_data = ""
        for line in data:
            elapsed = time.time() - start
            send_data += line

            if (len(send_data) > max_data) or (elapsed > max_time):
                result = sendit()
                send_data = ""
                start = time.time()
        if len(send_data):
            result = sendit()

        # Return the most recent JSON result we got back, or None if
        # we didn't make any requests.
        return result

    def stream_intervals(self, path, start = None, end = None):
        """
        Return a generator that yields each stream interval.
        Multiple requests are made to the server if the results
        get truncated.
        """
        params = {
            "path": path
        }
        if start is not None:
            params["start"] = repr(start)  # use repr to keep precision
        if end is not None:
            params["end"] = repr(end)

        while True:
            (intervals, restart) = self.curl.getjson("stream/intervals", params)
            for interval in intervals:
                yield interval
            if restart:
                # Restart where we left off
                params["start"] = repr(restart)
            else:
                break

    def stream_extract(self, path, start = None, end = None, bare = False):
        """
        Extract data from a stream.  Returns a generator that yields
        chunks of ASCII-formatted data that matches the database
        layout for the given path.  Multiple requests are made to the
        server if shorter requests get truncated.
        """
        params = {
            "path": path
        }
        if start is not None:
            params["start"] = repr(start)  # use repr to keep precision
        if end is not None:
            params["end"] = repr(end)
        params["bare"] = bare

        more = True
        while more:
            self.curl.get
            (intervals, more) = self.curl.getjson("stream/intervals", params)
            for interval in intervals:
                yield interval
            if more:
                # Restart where we left off
                params["start"] = repr(intervals[-1][1])