SHARE
TWEET

MongoDB Cache for Django

a guest Dec 29th, 2016 233 Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
  1. # -*- coding: utf-8 -*-
  2. #  Based on ---->
  3. #  Author:
  4. #  Karol Sikora <karol.sikora@laboratorium.ee>, (c) 2012
  5. #  Alireza Savand <alireza.savand@gmail.com>, (c) 2013, 2014, 2015
  6. #
  7. #
  8. # Remake by isox@vulners.com
  9. # Using TTL index collection
  10.  
  11. try:
  12.     import cPickle as pickle
  13. except ImportError:
  14.     import pickle
  15. import base64
  16. import re
  17. from datetime import datetime, timedelta
  18.  
  19. import pymongo
  20. from django.core.cache.backends.base import BaseCache
  21.  
  22.  
  23. class MongoDBCache(BaseCache):
  24.     def __init__(self, location, params):
  25.         BaseCache.__init__(self, params)
  26.         self.location = location
  27.         options = params.get('OPTIONS', {})
  28.         self._host = options.get('HOST', 'localhost')
  29.         self._port = options.get('PORT', 27017)
  30.         self._username = options.get('USERNAME')
  31.         self._password = options.get('PASSWORD')
  32.         self._database = options.get('DATABASE', 'django_cache')
  33.         self._collection = location
  34.         self._db, self._coll = self.initMongoConnection()
  35.  
  36.     def initMongoConnection(self):
  37.         if self._username is not None:
  38.             self.connection = pymongo.MongoClient(
  39.                 'mongodb://{0}:{1}@{2}:{3}/{4}'.format(self._username, self._password, self._host, self._port,
  40.                                                        self._database))
  41.         else:
  42.             self.connection = pymongo.MongoClient('mongodb://{0}:{1}/'.format(self._host, self._port))
  43.  
  44.         # Initialize key index
  45.         self.connection[self._database][self._collection].ensure_index('key', background=True)
  46.         # Initialize TTL index
  47.         # Elements will be deleted after 5 seconds expiration date will be passed
  48.         self.connection[self._database][self._collection].ensure_index('expires', background = True, expireAfterSeconds = 5)
  49.  
  50.         return self.connection[self._database], self.connection[self._database][self._collection]
  51.  
  52.  
  53.     def make_key(self, key, version=None):
  54.         """
  55.         Additional regexp to remove $ and . cachaters,
  56.        as they cause special behaviour in mongodb
  57.        """
  58.         key = super(MongoDBCache, self).make_key(key, version)
  59.  
  60.         return re.sub(r'\$|\.', '', key)
  61.  
  62.     def add(self, key, value, timeout=None, version=None):
  63.         """
  64.            Set a value in the cache if the key does not already exist. If
  65.            timeout is given, that timeout will be used for the key; otherwise
  66.            the default cache timeout will be used.
  67.  
  68.            Returns True if the value was stored, False otherwise.
  69.        """
  70.         key = self.make_key(key, version)
  71.         self.validate_key(key)
  72.  
  73.         return self._base_set('add', key, value, timeout)
  74.  
  75.     def set(self, key, value, timeout=None, version=None):
  76.         """
  77.            Set a value in the cache. If timeout is given, that timeout will be
  78.            used for the key; otherwise the default cache timeout will be used.
  79.        """
  80.         key = self.make_key(key, version)
  81.         self.validate_key(key)
  82.  
  83.         return self._base_set('set', key, value, timeout)
  84.  
  85.     def _base_set(self, mode, key, value, timeout=None):
  86.  
  87.         timeout = timeout or self.default_timeout
  88.  
  89.         # Only UTC here for Mongo auto-purge
  90.  
  91.         now = datetime.utcnow()
  92.         expires = now + timedelta(seconds=timeout)
  93.  
  94.         #
  95.  
  96.         pickled = pickle.dumps(value, pickle.HIGHEST_PROTOCOL)
  97.         encoded = base64.encodebytes(pickled).strip()
  98.  
  99.         if mode == 'set':
  100.             # Set sets new data. And there is no matter, that key exists
  101.             self._coll.update({'key':key} ,{'key':key, 'data': encoded, 'expires': expires}, upsert = True)
  102.  
  103.         elif mode == 'add':
  104.             if self._coll.find_one({'key': key}):
  105.                 return False
  106.             self._coll.insert({'key': key, 'data': encoded, 'expires': expires})
  107.  
  108.         return True
  109.  
  110.     def get(self, key, default=None, version=None):
  111.         """
  112.            Fetch a given key from the cache. If the key does not exist, return
  113.            default, which itself defaults to None.
  114.        """
  115.         key = self.make_key(key, version)
  116.         self.validate_key(key)
  117.  
  118.         data = self._coll.find_one({'key': key})
  119.         if not data:
  120.             return default
  121.  
  122.         unencoded = base64.decodebytes(data['data'])
  123.         unpickled = pickle.loads(unencoded)
  124.  
  125.         return unpickled
  126.  
  127.     def get_many(self, keys, version=None):
  128.         """
  129.            Fetch a bunch of keys from the cache. For certain backends (memcached,
  130.            pgsql) this can be *much* faster when fetching multiple values.
  131.  
  132.            Returns a dict mapping each key in keys to its value. If the given
  133.            key is missing, it will be missing from the response dict.
  134.        """
  135.         out = {}
  136.         parsed_keys = {}
  137.  
  138.         for key in keys:
  139.             pkey = self.make_key(key, version)
  140.             self.validate_key(pkey)
  141.             parsed_keys[pkey] = key
  142.  
  143.         data = self._coll.find({'key': {'$in': parsed_keys.keys()}})
  144.         for result in data:
  145.             unencoded = base64.decodebytes(result['data'])
  146.             unpickled = pickle.loads(unencoded)
  147.             out[parsed_keys[result['key']]] = unpickled
  148.  
  149.         return out
  150.  
  151.     def delete(self, key, version=None):
  152.         """
  153.            Delete a key from the cache, failing silently.
  154.        """
  155.         key = self.make_key(key, version)
  156.         self.validate_key(key)
  157.         self._coll.remove({'key': key})
  158.  
  159.     def has_key(self, key, version=None):
  160.         """
  161.            Returns True if the key is in the cache and has not expired.
  162.        """
  163.         key = self.make_key(key, version)
  164.         self.validate_key(key)
  165.         data = self._coll.find_one({'key': key})
  166.  
  167.         return data is not None
  168.  
  169.     def clear(self):
  170.         """Remove *all* values from the cache at once."""
  171.         self._coll.remove(None)
  172.  
  173.     def close(self, **kwargs):
  174.         """Close the cache connection"""
  175.         return self.connection.close()
RAW Paste Data
We use cookies for various purposes including analytics. By continuing to use Pastebin, you agree to our use of cookies as described in the Cookies Policy. OK, I Understand
 
Top