Recurse Center database server pairing task

data_store.py

# Python 3

from http.server import HTTPServer, BaseHTTPRequestHandler
from urllib.parse import urlparse, parse_qs


class DataStore(object):
    """A generic data storage object"""

    def __init__(self):
        """Initialise data store in memory"""
        self.store = {}

    def set(self, key, value):
        """Save the value to key (overwriting existing)"""
        self.store[key] = value

    def get(self, key):
        """Get the value for key and raise KeyError if unavailable"""
        return self.store[key]

    def __str__(self):
        return str(self.store)


class RequestHandler(BaseHTTPRequestHandler):
    """Handles requests"""

    def get(self, key):
        """Retrieve the data with key"""
        # key is a list: ?key=a&key=b --> ['a', 'b']
        try:
            list_of_data = []
            for k in key:
                list_of_data.append(data_store.get(bytes(k, 'utf-8')))
            data = b'\n'.join(list_of_data)
            self.send_response(200)
            self.send_header('Content-type', 'application/octet-stream')
            self.send_header('Content-length', len(data))
            self.end_headers()
            self.wfile.write(data)
        except KeyError:
            self.send_error(404, 'Key not found')

    def set(self, dictionary):
        """Save the data in dictionary (key, val)"""
        if len(dictionary) < 1:
            self.send_error(400, 'Invalid request (key is required)')
            return
        self.send_response(201)
        self.send_header('Content-type', 'text/plain')
        self.end_headers()
        for key in dictionary:
            # We only store the first occurrance
            data_store.set(bytes(key, 'utf-8'), bytes(dictionary[key][0], 'utf-8'))
            self.wfile.write(b'Content written for key ' + bytes(key, 'utf-8') + b'\n')

    def do_GET(self):
        """Serve a get request"""
        url = urlparse(self.path)
        path = url.path.strip('/').lower()
        query = parse_qs(url.query)
        if 'get' == path:
            try:
                self.get(query['key'])
            except KeyError:
                self.send_error(400, 'Invalid request (key is required)')
        elif 'set' == path:
            self.set(query)
        else:
            self.send_error(400, 'Invalid request (path must be /get or /set)')
        self.end_headers()

data_store = DataStore()

server = HTTPServer(('127.0.0.1', 4000), RequestHandler)
server.serve_forever()