Piston Redis OAuth Store

redis_store.py

import hashlib
import redis
from django.db import DEFAULT_DB_ALIAS
from django.db.models import Q, signals as db_signals
from piston.models import Consumer, Token
from piston.store import DataStore as PistonDataStore
from mirrio.api._redis import redis


redis = redis.Redis(host=settings.REDIS_HOST, port=settings.REDIS_PORT, db=settings.REDIS_DB)

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

def redis_model_cache(model_class):
  r"""
  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):
    r"""
    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.
    
    :arg q: A django `Q` object (to use on `model_class.objects.get()`)
    :arg model_class: A django model
    :arg expiry: When this key should expire from the cache
    :rtype: An instance of model_class
    """ % dict(model=model_class)

    # build an instance of model_class from dict:d
    def _builder(d):
      inst = model_class(**d)
      inst._state.adding = False # so django doesn't try to overwrite 
      inst._state.db = using or DEFAULT_DB_ALIAS
      return inst
    
    # save an instance of model_class
    def _cache_model(key, obj):
      d = {}
      for field in obj._meta.fields:
        if field.get_internal_type() == 'FileField':
          continue
        d[field.attname] = getattr(obj, field.attname)
      # save the object
      redis.hmset(key, d)
      redis.expire(key, expiry)
      return d
    
    # 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
    pk_key = 'q' + md5(model_class.__name__ + str(q))
    # see if this query has been performed before
    pk = redis.get(pk_key)
    if pk is not None:
      # HEY WE FOUND A PK FOR THIS QUERY LETS TRY TO GET IT FROM REDIS
      key = 'pk' + model_class.__name__ + pk
      try:
        #print 'cache hit key', key
        return _builder(redis.hgetall(key))
      except:
        redis.delete(key)
        redis.delete(pk_key)
    else:
      # one caveat to doing it this way is that if we have a cache miss on retrieving
      # the query=>pk, we have to fetch the pk of the matching object, and might as well
      # get the entire thing
      obj = model_class._default_manager.using(using).get(q)
      # save the query => pk cache
      redis.set(pk_key, str(obj.pk))
      redis.expire(pk_key, expiry)
      
      # now we do normal row caching
      key = 'pk' + model_class.__name__ + str(obj.pk)
      #print 'cache miss key', key
      if not redis.exists(key): # but don't re-cache if it's not necessary
        #print 'caching key', key
        _cache_model(key, obj)

      return obj

  if not hasattr(model_class, 'redis_cache_cached'): # only connect these things once
    def _clear(sender, instance, *args, **kwargs):
      key = 'pk' + model_class.__name__ + str(instance.pk)
      #print 'expiring key', key
      redis.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, 'redis_cache_cached', True)
  
  return _cacher
  
consumer_cache = redis_model_cache(Consumer)
token_cache = redis_model_cache(Token)


class RedisDataStore(PistonDataStore):
  def lookup_consumer(self, key):
    self.consumer = consumer_cache(Q(key=key))
    self.consumer.key = self.consumer.key.encode('ascii')
    self.consumer.secret = self.consumer.secret.encode('ascii')
    return self.consumer

  def lookup_token(self, token_type, token):
    if token_type == 'request':
      token_type = Token.REQUEST
    elif token_type == 'access':
      token_type = Token.ACCESS
    try:
      self.request_token = token_cache(Q(token_type=token_type, key=token))
      return self.request_token
    except Token.DoesNotExist:
      return None
    
  def lookup_nonce(self, oauth_consumer, oauth_token, nonce):
    key = 'nonce' + md5(oauth_consumer.key + oauth_consumer.secret 
                        + getattr(oauth_token, 'key', '')
                        + getattr(oauth_token, 'secret', '') + nonce)
    r = redis.get(key)
    if not r:
      redis.set(key, 'v')
      redis.expire(key, 60*60*24)
      return None
    return r