Key/value ORM storage

# -*- coding: utf-8 -*-

"""
A key/value store based on a relational database table.

Usage::

    # Store a value, retrieve, show and delete a value for a key.
    KeyValue.set(u'foo', u'bar')
    print KeyValue.get(u'foo')
    KeyValue.delete(u'foo')

    # Increment a counter (pre-initialized at zero) and show the result.
    hits = KeyValue.incr(u'hits')
    print 'Hits:', hits

    # Use optional JSON serialization (`None`, boolean, number, string,
    # and list types only).
    KeyValue.set(u'fibonacci', [0, 1, 1, 2, 3, 5, 8], True)
    print ' -> '.join(KeyValue.get(u'fibonacci', True))
"""

import simplejson as json
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import scoped_session, sessionmaker
from sqlalchemy.schema import Column, MetaData
from sqlalchemy.types import Unicode, UnicodeText


# Prepare metadata, session, and Base class with query attribute.
metadata = MetaData()
dbsession = scoped_session(sessionmaker())
Base = declarative_base(metadata=metadata)
Base.q = dbsession.query_property()


class KeyValue(Base):
    """A key/value record."""
    __tablename__ = 'key_values'

    key = Column(Unicode, primary_key=True)
    value = Column(UnicodeText)

    @classmethod
    def get(cls, key, serialized=False):
        """Retrieve the value stored at the given key."""
        kv = cls.q.get(key)
        if kv is None:
            # Key does not exist.
            return
        if serialized:
            return json.loads(kv.value)
        return kv.value

    @classmethod
    def set(cls, key, value, serialized=False):
        """Store the given value at the given key.

        If the key exists, the value is overwritten, otherwise the key is
        created and the value assigned to it.
        """
        if serialized:
            value = unicode(json.dumps(value))
        kv = cls.q.get(key)
        if kv is None:
            dbsession.add(KeyValue(key=key, value=value))
        else:
            kv.value = value
        dbsession.commit()

    @classmethod
    def incr(cls, key, amount=1):
        """Increment an integer value by the given amount, save and return the
        resulting number.
        """
        count = cls.get(key)
        if count is None:
            count = 0
        else:
            try:
                count = int(count)
            except ValueError:
                raise ValueError('Key exists, but has a non-integer value.')
        count += amount
        cls.set(key, unicode(count))
        return count

    @classmethod
    def delete(cls, key):
        """Delete the key and its value."""
        kv = cls.q.get(key)
        if kv is not None:
            dbsession.delete(kv)
            dbsession.commit()