A Redis-like in-memory database

simple_db.py

#!/usr/bin/env python
"""
This is my solution to the Simple Database problem posted on:

http://www.thumbtack.com/challenges/software-engineer

Candidate: David Morton
Email:     dcmorton@gmail.com

Usage:

./simple_db.py
"""
import sys
from collections import defaultdict


class Database(object):
    """
    An simple in-memory database that allows values to
    be set and accessed via the following operations:

    SET <name> <value>
    GET <name>
    UNSET <name>
    NUMEQUALTO <value>
    END

    Transactional functionality is provided via the following
    commands:

    BEGIN
    ROLLBACK
    COMMIT
    """
    def __init__(self):
        """
        Initialize the In-Memory Database

        The following variables are bound to the Database instance:

        counts: A mapping between values and the count of keys
                for which the corresponding value is stored

        values: A mapping between keys and the list of values;
                we're using a list so that we can revert to previous
                values in the event of transaction rollbacks

        transaction_blocks: A list of sets that correspond to
                            nested transactions. Each set contains
                            the keys that have been altered within
                            the respective transaction
        """
        self.counts = defaultdict(int)
        self.values = defaultdict(list)
        self.transaction_blocks = []

    def _commit(self):
        """
        Close all open transaction blocks, permanently applying
        the changes made in them. Print nothing if successful, or
        print NO TRANSACTION if no transaction is in progress.

        The only keys in self.values that are modified are those
        in the in superset of self.transaction_blocks.
        """
        if not self.transaction_blocks:
            print 'NO TRANSACTION'
        else:
            # Get the superset of the keys in all transaction block sets
            transaction_keys = set()
            for block in self.transaction_blocks:
                transaction_keys |= set(block)
            for key in transaction_keys:
                # If the key has been 'unset', don't retain it
                if self._get_value(key) is None:
                    self.values[key].pop()
                # Otherwise self.values[key] should contain the
                # the element as it was set last.
                else:
                    self.values[key] = [self._get_value(key)]
            # Reset the list of containing transaction blocks
            self.transaction_blocks = []

    def _get(self, key):
        """
        Print out the value that has been stored for the key
        or NULL if the key is unset.
        """
        value = self._get_value(key)
        # None is indicates UNSET keys
        print value if value is not None else 'NULL'

    def _get_value(self, key):
        """
        A utility method used to get the value stored for
        a given key. Return None if the key is unset.
        """
        if self.values[key]:
            return self.values[key][-1]

    def _rollback(self):
        """
        Undo all of the commands issues in the most recent transaction
        block, and close the block. Print nothing if successful, or
        print NO TRANSACTION if no transaction is in progress.
        """
        if not self.transaction_blocks:
            print 'NO TRANSACTION'
        else:
            # Iterate over all of the keys in the transaction block
            for key in self.transaction_blocks.pop():
                # Pop the the value that was (un)set within the
                # current transaction block and update the counts.
                old_value = self.values[key].pop()
                self._update_counts(old_value, self._get_value(key))

    def _set(self, key, value):
        """
        Set the key name to the value

        This needs to be handled differently based on the following criteria:

        1) Whether or not the SET command is issued within a transaction block
        2) Whether or not there has been a previous SET command issued
           within the transaction block
        3) Whether or not the key already has a value (regardless of whether
           or not the SET command is issued within a transaction block)
        """
        # Check to see if the key in the current transaction block
        not_in_block = self.transaction_blocks and key not in self.transaction_blocks[-1]

        # Update the counts for the potential change in value for the given key
        self._update_counts(self._get_value(key), value)

        # Append the value to the values list if it doesn't
        # exist in values or in the current transaction block
        if not self.values[key] or not_in_block:
            self.values[key].append(value)
            if self.transaction_blocks:
                self.transaction_blocks[-1].add(key)
        # If the value has been set and you're not in a transaction block
        # or if you're setting the value multiple times in the same
        # transaction block, update the last element in self.values[key]
        else:
            self.values[key][-1] = value

    def _unset(self, key):
        """
        Unset the value of key.
        """
        # I'm assuming that, if a user issues an 'UNSET'
        # for a key that hasn't been set, it's a no-op
        if not self.values[key]:
            pass
        # If you're in a transaction block, set a value of
        # None for purposes of committing and rolling back
        elif self.transaction_blocks:
            self._set(key, None)
        # If you're not in a transaction block, decrement
        # the counter for the key, and pop the key
        else:
            self._update_counts(self._get_value(key), None)
            self.values.pop(key)

    def _update_counts(self, old_value, new_value):
        """
        Decrement the counter for the old_value and
        augment the counter for the new value.
        """
        if new_value is not None:
            self.counts[new_value] += 1
        if old_value is not None:
            self.counts[old_value] -= 1
            # Remove the reference in self.counts if it's no longer needed
            if self.counts[old_value] == 0:
                self.counts.pop(old_value)

    def parse_command(self, line):
        """
        Parse the command from stdin and perform the appropriate action
        """
        cmd = line.split()
        if cmd:  # Ignore blank lines
            func = cmd[0].upper()  # Accept lowercase commands
            if func == 'BEGIN':
                self.transaction_blocks.append(set())
            elif func == 'COMMIT':
                self._commit()
            elif func == 'END':
                sys.exit(0)  # The 'END' command terminates the session
            elif func == 'GET':
                self._get(cmd[1])
            elif func == 'NUMEQUALTO':
                print self.counts[cmd[1]]
            elif func == 'ROLLBACK':
                self._rollback()
            elif func == 'SET':
                self._set(*cmd[1:])
            elif func == 'UNSET':
                self._unset(cmd[1])
            else:
                print 'Invalid Command: "{}"'.format(func)


def main():
    database = Database()
    while True:
        try:
            database.parse_command(raw_input())
        except KeyboardInterrupt:
            sys.exit(0)

if __name__ == '__main__':
    main()