A Python decorator that caches the result of a function on MongoDB

mongo_cache.py

# -*- coding: utf-8 -*-

import base64
import cPickle as pickle
import hashlib
from functools import wraps


class NoopSerializer(object):
    def serialize(self, obj):
        return obj

    def deserialize(self, serialized):
        return serialized


class PickleSerializer(object):
    def serialize(self, obj):
        return base64.b64encode(pickle.dumps(obj, protocol=-1))

    def deserialize(self, serialized):
        return pickle.loads(base64.b64decode(serialized))


def mongo_cache(db_conn, prefix='cache_', capped=True, capped_size=1000000000,
                hash_keys=True, serializer=PickleSerializer):
    def decorator(func):
        serializer_ins = serializer()
        col_name = '%s%s' % (prefix, func.__name__)
        if capped:
            try:
                db_conn.create_collection(col_name, capped=capped, size=capped_size)
            except:
                pass

        cache_col = db_conn[col_name]
        cache_col.ensure_index('key', unique=True)

        @wraps(func)
        def wrapped_func(*args, **kwargs):
            cache_key = pickle.dumps((args, kwargs), protocol=-1)
            if hash_keys:
                cache_key = hashlib.md5(cache_key).hexdigest()
            else:
                cache_key = base64.b64encode(cache_key)

            cached_obj = cache_col.find_one(dict(key=cache_key))
            if cached_obj:
                return serializer_ins.deserialize(cached_obj['result'])

            ret = func(*args, **kwargs)
            cache_col.update(
                {'key': cache_key},
                {'$set': {'result': serializer_ins.serialize(ret)}},
                upsert=True
            )

            return ret

        return wrapped_func

    return decorator