Riak session backend for Django

riak.py

from __future__ import absolute_import

from datetime import datetime, timedelta

import riak

from django.conf import settings
from django.contrib.sessions.backends.base import SessionBase, CreateError


class SessionStore(SessionBase):
    def __init__(self, session_key=None):
        self._client = riak.RiakClient(port=8087,
                                       transport_class=riak.RiakPbcTransport)
        self._bucket = self._client.bucket('sessions')
        super(SessionStore, self).__init__(session_key)

    def load(self):
        session = self._bucket.get(self.session_key)

        # the Riak client always returns an object, we have to check
        # exists() to see whether or not it's "fake"
        if session.exists():
            session_data = session.get_data()

            # only return unexpired sessions
            expire_date = datetime.fromtimestamp(session_data['_expire_date'])
            now = datetime.now()
            if (now - expire_date) < timedelta(seconds=settings.SESSION_COOKIE_AGE):
                decoded = self.decode(session_data['_encoded_data'])
                return decoded

        self.create()
        return {}

    def create(self):
        while True:
            self.session_key = self._get_new_session_key()
            try:
                self.save(must_create=True)
            except CreateError:
                # the key wasn't unique, try again
                continue
            self.modified = True
            self._session_cache = {}
            return

    def save(self, must_create=False):
        # without true single-node consistency the best we can do is a
        # get to ensure the value doesn't already exist in the DB
        # it is assumed that your R+W > N
        if must_create:
            current_value = self._bucket.get(self._session_key)
            if current_value.exists():
                return CreateError

        encoded_data = self.encode(self._get_session(no_load=must_create))
        expire_date = int(self.get_expiry_date().strftime("%s"))

        session = self._bucket.new(self.session_key, data={
            '_encoded_data': encoded_data,
            '_expire_date': expire_date,
        })
        session.store()

    def exists(self, session_key):
        session = self._bucket.get(session_key)
        return session.exists()

    def delete(self, session_key=None):
        if session_key is None:
            if self._session_key is None:
                return
            session_key = self._session_key

        session = self._bucket.get(session_key)
        if session.exists():
            session.delete()

NOTES:

• Based on https://bitbucket.org/stodge/django-riak-session
• Fixes a handful of bugs or oversights

TODO:

• Right now it assumes it's connecting to Riak on the localhost. I'm looking into adding server-pools to the Python client now.