Source

pythonhk / src / pythonhk / cp / session.py

import cPickle as pickle
import logging
import threading
import time

from pprint import pformat

from cherrypy.lib import sessions
from redis import Redis as _RedisClient


logger = logging.getLogger(__name__)


class RedisSession(sessions.Session):

    locks = {}

    prefix = "cp-session:"

    @classmethod
    def setup(cls, **kwargs):
        for k, v in kwargs.items():
            setattr(cls, k, v)
        cls.cache = cache = _RedisClient(**kwargs)
        redis_info = cache.info()
        logger.info("Redis server ready.\n%s" % pformat(redis_info))

    def _exists(self):
        return self.cache.exists(self.prefix + self.id)

    def _load(self):
        data = self.cache.get(self.prefix + self.id)
        if data:
            retval = pickle.loads(data)
            return retval

    def _save(self, expiration_time):
        key = self.prefix + self.id
        td = int(time.mktime(expiration_time.timetuple()))

        data = pickle.dumps((self._data, expiration_time),
                            pickle.HIGHEST_PROTOCOL)

        def critical_section(pipe):
            pipe.multi()
            pipe.set(key, data)
            pipe.expireat(key, td)

        self.cache.transaction(critical_section, key)

    def _delete(self):
        self.cache.delete(self.prefix + self.id)

    def acquire_lock(self):
        self.locked = True
        self.locks.setdefault(self.id, threading.RLock()).acquire()

    def release_lock(self):
        self.locks[self.id].release()
        self.locked = False

    def __len__(self):
        """Return the number of active sessions."""
        keys = self.cache.keys(self.prefix + '*')
        return len(keys)