from django.db.models.manager import Manager
from django.db.models.query import QuerySet
from django.core.cache import cache

DEFAULT_CACHE_TIME = 15*60 # 15 minutes
# TODO
# - Come up with a better method for invalidation
# - Add invalidation for count() when a queryset is invalidated
# - Find a way to make AutoCacheManager override `objects` in models
# - Add some handling to allow CacheManager to react differently based on query type (get, count, filter, select_related)

# CacheManager -- A manager to store and retrieve cached objects using CACHE_BACKEND
# (Optional) <string key_prefix> -- the key prefix for all cached objects on this model [default: db_table]
# (Optional) <int timeout> -- in seconds, the maximum time before data is invalidated

# cachemanager.clean() -- Invalidates cached data
# <instance/queryset data> -- the queryset, or instance of the object to be invalidated
class CacheManager(Manager):
def __init__(self, *args, **kwargs):
self.key_prefix = kwargs.pop('key_prefix', None)
self.timeout = kwargs.pop('timeout', None)
Manager.__init__(self)

def get_query_set(self):
return CachedQuerySet(self.model, self.timeout, self.key_prefix)

# clean will accept either a queryset or an instance
def clean(self, data):
# invalidate the .get() request
if isinstance(data, self.model):
self.clean(self.filter(pk=self._get_pk_val()))
elif isinstance(data, CachedQuerySet):
self.clean()
else:
raise TypeError("instance or queryset required for data, got %r" % (data,))

# CachedQuerySet -- Extends the QuerySet object -- additionally adds a .cache() method

# queryset.cache() -- Overrides CacheManager's options for this QuerySet
# (Optional) <string key_prefix> -- the key prefix for all cached objects on this model [default: db_table]
# (Optional) <int timeout> -- in seconds, the maximum time before data is invalidated

# queryset.clean() -- Removes queryset from the cache -- recommended to use cachemanager.clean()
# must be called as the last method of the queryset
# <instance/queryset data> -- the queryset, or instance of the object to be invalidated
class CachedQuerySet(QuerySet):
def __init__(self, model=None, *args, **kwargs):
self.__cache_key = None
self.key_prefix = kwargs.pop('key_prefix', model and model._meta.db_table or '')
self.timeout = kwargs.pop('timeout', getattr(cache, 'default_timeout', DEFAULT_CACHE_TIME))
if not isinstance(self.key_prefix, basestring):
raise TypeError("string required for key_prefix, got %r" % (self.key_prefix,))
if not isinstance(self.timeout, int):
raise TypeError("integer required for timeout, got %r" % (self.timeout,))
QuerySet.__init__(self, model)

def _get_sorted_clause_key(self):
return (isinstance(i, basestring) and i.lower().replace('`', '').replace("'", '') or str(tuple(sorted(i))) for i in self._get_sql_clause())

def _get_cache_key(self, xtra=''):
if not self.__cache_key:
self.__cache_key = self.key_prefix + str(hash(''.join(self._get_sorted_clause_key()))) + xtra
return self.__cache_key

def _get_data(self):
data = cache.get(self._get_cache_key())
if data is None:
data = QuerySet._get_data(self)
cache.set(self._get_cache_key(), [d for d in data], self.timeout)
return data

def count(self):
count = cache.get(self._get_cache_key('count'))
if count is None:
count = int(QuerySet.count(self))
cache.set(self._get_cache_key('count'), count, self.timeout)
return count

def cache(self, *args, **kwargs):
self.key_prefix = kwargs.pop('key_prefix', self.key_prefix)
self.timeout = kwargs.pop('timeout', self.timeout)
if not isinstance(self.key_prefix, basestring):
raise TypeError("string required for key_prefix, got %r" % (self.key_prefix,))
if not isinstance(self.timeout, int):
raise TypeError("integer required for timeout, got %r" % (self.timeout,))

def clean(self):
cache.delete(self._get_cache_key())

# AutoCacheManager -- automatically replaces objects with a CacheManager and
# overrides save/delete functions to update the cache as needed
class AutoCacheManager(object):
# TODO: find a way to make `objects` actually set correctly
objects = CacheManager()
no_cache = Manager()

# TODO: know how to clean this w/o guessing for CacheManager
def delete(self, *args, **kwargs):
self.__class__.objects.clean(self)
Manager.delete(self, *args, **kwargs)

def save(self, *args, **kwargs):
Manager.save(self, *args, **kwargs)
self.__class__.objects.clean(self)