Django Model Cache and Cached Piston OAuth Store

model_cache.py

# The author disclaims copyright to this source code.  In place of a legal 
# notice, here is a blessing:
#
#    May you do good and not evil.
#    May you find forgiveness for yourself and forgive others.
#    May you share freely, never taking more than you give.

import hashlib, cPickle as pickle, logging
from django.db.models import signals as db_signals
from django.core.cache import cache

logger = logging.getLogger('caching')
def debug(*args):
    logger.debug('cache debug {}'.format(' '.join(args)))

md5 = lambda v: hashlib.md5(v).hexdigest()

def ModelCache(model_class, verbose=False, exc_on_not_found=True):
    """
    Returns a callable which will cache the model passed in. See the docstring 
    for the returned callable.

    :arg model_class: a django.model.Model class
    :rtype: a callable which can be used to access and cache rows from 
            the model class
    """
    def _cacher(q, expiry=60*60*24, using=None):
        """
        This is a cache of %(model)s. You can retrieve single objects using a 
        django.models.Q object. The method will cache which exact row is 
        identified by this query and cache a it's primary key, in addition to 
        the row itself.

        Use it like this::

           from django.contrib.auth.user import User
           from django.models import Q
           from model_cache import ModelCache

           user_cache = ModelCache(User)
           
           user = user_cache(Q(username='samuraisam', is_active=True)) # woo!
        
        :arg q: A django `Q` object (to use on `model_class.objects.get()`)
        :arg model_class: A django model class
        :arg expiry: When this key should expire from the cache (in secs)
        :arg using: Tells Django which database to use when quering for the obj
        :rtype: An instance of model_class
        """ % dict(model=model_class.__name__)

        # build an instance of model_class from dict:d
        def _builder(d):
            return pickle.loads(d)
        
        # save an instance of model_class
        def _cache_model(key, obj):
            cache.set(key, pickle.dumps(obj), expiry)
        
        # we save a hash of the query and save it to the pk it actually 
        # represents this way we can make cache arbitrary queries that lookup
        # the same object. we also include the model's class name, and
        # a hash of the entire class dictionary (in case the contents change
        # in some way, we'll be prepared)
        mh = md5('{}{}{}'.format(str(q), model_class.__name__, 
                                 str(model_class.__dict__)))
        mk = md5('{}{}'.format(model_class.__name__, str(model_class.__dict__)))
        pk_key = 'q{}'.format(mh)
        # see if this query has been performed before
        pk = cache.get(pk_key)
        obj = None
        if pk is not None:
            # HEY WE FOUND A PK FOR THIS QUERY LETS TRY TO GET IT FROM CACHE
            key = 'pk{}{}'.format(mk, pk)
            try:
                if verbose:
                    debug('cache hit key', key)
                obj = _builder(cache.get(key))
                obj._from_cache = True
                return obj
            except Exception, e:
                if verbose:
                    debug('cache build error', str(e))
                cache.delete(key)
                cache.delete(pk_key)

        if obj is None:
            # DERP, CACHE MISS
            try:
                obj = model_class._default_manager.using(using).get(q)
            except model_class.DoesNotExist:
                if exc_on_not_found:
                    raise
                else:
                    return None
            # save the query => pk cache
            cache.set(pk_key, str(obj.pk), expiry)
            
            # now we do normal row caching
            key = 'pk{}{}'.format(mk, str(obj.pk))
            if verbose:
                debug('cache miss key', key)

            # but don't re-cache if it's not necessary
            if not cache.has_key(key):
                if verbose:
                    debug('caching key', key)
                _cache_model(key, obj)
            obj._from_cache = False
            return obj

    # only connect these things once
    if not hasattr(model_class, '_model_cached'): 
        def _clear(sender, instance, *args, **kwargs):
            mk = md5('{}{}'.format(model_class.__name__, 
                                   str(model_class.__dict__)))
            key = 'pk{}{}'.format(mk, str(instance.pk))
            if verbose:
                debug('expiring key', key)
            cache.delete(key)

        db_signals.post_save.connect(_clear, sender=model_class, weak=False)
        db_signals.post_delete.connect(_clear, sender=model_class, weak=False)

        setattr(model_class, '_model_cached', True)
    
    return _cacher

model_cache_test.py

# The author disclaims copyright to this source code.  In place of a legal 
# notice, here is a blessing:
#
#    May you do good and not evil.
#    May you find forgiveness for yourself and forgive others.
#    May you share freely, never taking more than you give.

from time import sleep
from unittest import TestCase
from django.contrib.auth.models import User
from django.db.models import Q
from model_cache import ModelCache


class ModelCacheTest(TestCase):
    def setUp(self):
        self._user = User.objects.create_user(
            username='samuraisam', email='[email protected]', password='yourmother'
        )
        self._user.is_active = True
        self._user.save()

    def tearDown(self):
        self._user.delete()

    def test_usermodel_cache_works(self):
        user_cache = ModelCache(User, verbose=True)

        # cold cache
        u = user_cache(Q(username='samuraisam'))
        self.assertEquals(u._from_cache, False)
        self.assertEquals(u.is_superuser, False)

        # should be cached now
        u = user_cache(Q(username='samuraisam'))
        self.assertEquals(u._from_cache, True)
        self.assertEquals(u.is_superuser, False)

        u.is_superuser = True
        u.save()

        # should be expired
        u = user_cache(Q(username='samuraisam'))
        self.assertEquals(u._from_cache, False)
        self.assertEquals(u.is_superuser, True)

        # and again warm
        u = user_cache(Q(username='samuraisam'))
        self.assertEquals(u._from_cache, True)
        self.assertEquals(u.is_superuser, True)

        user_cache_fast_expire = ModelCache(User, verbose=True)
        u = User.objects.create_user(username='samuraisam1', 
                                     password='yourmother1', 
                                     email='[email protected]')
        u.is_active = True
        u.save()

        u = user_cache_fast_expire(Q(username='samuraisam1'), expiry=1)
        self.assertEquals(u._from_cache, False)
        u = user_cache_fast_expire(Q(username='samuraisam1'), expiry=1)
        self.assertEquals(u._from_cache, True)

        sleep(2)

        u = user_cache_fast_expire(Q(username='samuraisam1'), expiry=1)
        self.assertEquals(u._from_cache, False)

    def test_dont_raise(self):
        user_cache = ModelCache(User, verbose=True, exc_on_not_found=False)
        self.assertEquals(user_cache(Q(username='omg')), None)
        u = User.objects.create_user(username='omg',password='omg',email='[email protected]')
        u.is_active = True
        u.save()
        self.assertEquals(user_cache(Q(username='omg')).pk, u.pk)
        u.delete()
        self.assertEquals(user_cache(Q(username='omg')), None)

    def test_raise(self):
        user_cache = ModelCache(User, verbose=True)
        with self.assertRaises(User.DoesNotExist):
            user_cache(Q(username='omg'))
        u = User.objects.create_user(username='omg',password='omg',email='[email protected]')
        u.is_active = True
        u.save()
        self.assertEquals(user_cache(Q(username='omg')).pk, u.pk)
        u.delete()
        with self.assertRaises(User.DoesNotExist):
            self.assertEquals(user_cache(Q(username='omg')), None)