dbserver.py

#!/usr/bin/env python
"""
Simple DB server (key value store) that works over HTTP on localhost.
  - GET /set?<key>=<value>  Sets <key> to <value>
  - GET /get?key=<key>      Gets the value for key <key>

Usage:
  dbserver.py [--port=<int>]

Options:
  -h --help         Show this screen.
  -p --port=<int>   Port to listen on [default: 4000].
"""
from BaseHTTPServer import BaseHTTPRequestHandler, HTTPServer
from docopt import docopt
import urlparse
import os


DATA = {}                             # this servers as the in-memory data store
DATA_DIR = './data'

class DBRequestHandlerBase(BaseHTTPRequestHandler):
    """
    Handles DB requests using storage-subclass `get` and `set` methods.
    """
    def do_GET(self):
        """Base response for all GET requests."""
        parsed_path = urlparse.urlparse(self.path)
        path = parsed_path.path
        #
        # 1. validate path
        if path not in ['/get', '/set']:
            self.send_error(500, 'Error: use /set?<key>=<val> or /get?key=<key>')
            return
        #
        # 2. validate query string
        query_dict = urlparse.parse_qs(parsed_path.query)
        if len(query_dict) != 1:
            self.send_error(500, 'Error: query must have a single field=value pair')
            return
        query_item = query_dict.items()[0] 
        q_field, q_value = query_item[0], query_item[1][0]
        #
        # 3. process request
        if path == '/set':
            self.set(q_field, q_value)                 # does actual write to DB
            self.respond(http_code=200, http_body='')
        elif path == '/get':
            if q_field != 'key':
                self.send_error(500, 'Error: use /get?key=<keyname> format')
                return
            value = self.get(q_value)                 # does actual read from DB
            if value is None:
                self.respond(http_code=404, http_body='')
                return
            self.respond(http_code=200, http_body=value)

    def respond(self, http_code=200, http_body=''):
        """Respond with default headers the `http_code` and `http_body` given."""
        self.send_response(http_code)
        self.send_header('Content-type', 'text/plain')
        self.end_headers()
        self.wfile.write(http_body)


class InMemoryDBRequestHandler(DBRequestHandlerBase):
    """
    Uses global dictionary `DATA` as data store for DB functionality.
    """
    def get(self, key):
        """Gets key `key` from DB. Returns None if `key` is not present."""
        return DATA.get(key, None)
    
    def set(self, key, value):
        """Sets key `key` to value `value` in DB."""
        DATA[key] = value


class PersistedDBRequestHandler(DBRequestHandlerBase):
    """
    Store data to disk as files in `DATA_DIR`; uses `DATA` dict as cache.
    """
    def get(self, key):
        """Gets key `key` from DB. Returns None if `key` is not present."""
        return DATA.get(key, None)
        # assumption: contents of `DATA` are always fresh
    
    def set(self, key, value):
        """Sets key `key` to value `value` in DB."""
        fpath = os.path.join(DATA_DIR, key)
        with open(fpath, 'w') as f:
            f.write(value)
        DATA[key] = value


class PersistedDBServer(HTTPServer):
    """
    Custom HTTPServer that pre-loads data from `DATA_DIR` into `DATA` dict.
    """
    def __init__(self, *args, **kwargs):
        HTTPServer.__init__(self, *args, **kwargs)       # old-style inheritence
        # ensure `DATA_DIR` exists
        if not os.path.exists(DATA_DIR):
            os.mkdir(DATA_DIR)
        # load data from files in `DATA_DIR` into `DATA` in-memory dictionary
        for fname in os.listdir(DATA_DIR):
            fpath = os.path.join(DATA_DIR, fname)
            with open(fpath, 'r') as f:
                DATA[fname] = f.read()


def run_server(port):
    """
    The main run loop for the DB server.
    """
    server_opts = ('localhost', port)
    server = PersistedDBServer(server_opts, PersistedDBRequestHandler)
    print 'Started Simple DB server on localhost port', port
    try:
        server.serve_forever()
    except KeyboardInterrupt:
        server.server_close()
        print 'Stopped Simple DB server'


if __name__ == '__main__':
    arguments = docopt(__doc__, version='SimpleDB v0.1')
    port_int = int(arguments['--port'])
    run_server(port=port_int)