SickRage/lib/shove/__init__.py

520 lines
15 KiB
Python

# -*- coding: utf-8 -*-
'''Common object storage frontend.'''
import os
import zlib
import urllib
try:
import cPickle as pickle
except ImportError:
import pickle
from collections import deque
try:
# Import store and cache entry points if setuptools installed
import pkg_resources
stores = dict((_store.name, _store) for _store in
pkg_resources.iter_entry_points('shove.stores'))
caches = dict((_cache.name, _cache) for _cache in
pkg_resources.iter_entry_points('shove.caches'))
# Pass if nothing loaded
if not stores and not caches:
raise ImportError()
except ImportError:
# Static store backend registry
stores = dict(
bsddb='shove.store.bsdb:BsdStore',
cassandra='shove.store.cassandra:CassandraStore',
dbm='shove.store.dbm:DbmStore',
durus='shove.store.durusdb:DurusStore',
file='shove.store.file:FileStore',
firebird='shove.store.db:DbStore',
ftp='shove.store.ftp:FtpStore',
hdf5='shove.store.hdf5:HDF5Store',
leveldb='shove.store.leveldbstore:LevelDBStore',
memory='shove.store.memory:MemoryStore',
mssql='shove.store.db:DbStore',
mysql='shove.store.db:DbStore',
oracle='shove.store.db:DbStore',
postgres='shove.store.db:DbStore',
redis='shove.store.redisdb:RedisStore',
s3='shove.store.s3:S3Store',
simple='shove.store.simple:SimpleStore',
sqlite='shove.store.db:DbStore',
svn='shove.store.svn:SvnStore',
zodb='shove.store.zodb:ZodbStore',
)
# Static cache backend registry
caches = dict(
bsddb='shove.cache.bsdb:BsdCache',
file='shove.cache.file:FileCache',
filelru='shove.cache.filelru:FileLRUCache',
firebird='shove.cache.db:DbCache',
memcache='shove.cache.memcached:MemCached',
memlru='shove.cache.memlru:MemoryLRUCache',
memory='shove.cache.memory:MemoryCache',
mssql='shove.cache.db:DbCache',
mysql='shove.cache.db:DbCache',
oracle='shove.cache.db:DbCache',
postgres='shove.cache.db:DbCache',
redis='shove.cache.redisdb:RedisCache',
simple='shove.cache.simple:SimpleCache',
simplelru='shove.cache.simplelru:SimpleLRUCache',
sqlite='shove.cache.db:DbCache',
)
def getbackend(uri, engines, **kw):
'''
Loads the right backend based on a URI.
@param uri Instance or name string
@param engines A dictionary of scheme/class pairs
'''
if isinstance(uri, basestring):
mod = engines[uri.split('://', 1)[0]]
# Load module if setuptools not present
if isinstance(mod, basestring):
# Isolate classname from dot path
module, klass = mod.split(':')
# Load module
mod = getattr(__import__(module, '', '', ['']), klass)
# Load appropriate class from setuptools entry point
else:
mod = mod.load()
# Return instance
return mod(uri, **kw)
# No-op for existing instances
return uri
def synchronized(func):
'''
Decorator to lock and unlock a method (Phillip J. Eby).
@param func Method to decorate
'''
def wrapper(self, *__args, **__kw):
self._lock.acquire()
try:
return func(self, *__args, **__kw)
finally:
self._lock.release()
wrapper.__name__ = func.__name__
wrapper.__dict__ = func.__dict__
wrapper.__doc__ = func.__doc__
return wrapper
class Base(object):
'''Base Mapping class.'''
def __init__(self, engine, **kw):
'''
@keyword compress True, False, or an integer compression level (1-9).
'''
self._compress = kw.get('compress', False)
self._protocol = kw.get('protocol', pickle.HIGHEST_PROTOCOL)
def __getitem__(self, key):
raise NotImplementedError()
def __setitem__(self, key, value):
raise NotImplementedError()
def __delitem__(self, key):
raise NotImplementedError()
def __contains__(self, key):
try:
value = self[key]
except KeyError:
return False
return True
def get(self, key, default=None):
'''
Fetch a given key from the mapping. If the key does not exist,
return the default.
@param key Keyword of item in mapping.
@param default Default value (default: None)
'''
try:
return self[key]
except KeyError:
return default
def dumps(self, value):
'''Optionally serializes and compresses an object.'''
# Serialize everything but ASCII strings
value = pickle.dumps(value, protocol=self._protocol)
if self._compress:
level = 9 if self._compress is True else self._compress
value = zlib.compress(value, level)
return value
def loads(self, value):
'''Deserializes and optionally decompresses an object.'''
if self._compress:
try:
value = zlib.decompress(value)
except zlib.error:
pass
value = pickle.loads(value)
return value
class BaseStore(Base):
'''Base Store class (based on UserDict.DictMixin).'''
def __init__(self, engine, **kw):
super(BaseStore, self).__init__(engine, **kw)
self._store = None
def __cmp__(self, other):
if other is None:
return False
if isinstance(other, BaseStore):
return cmp(dict(self.iteritems()), dict(other.iteritems()))
def __del__(self):
# __init__ didn't succeed, so don't bother closing
if not hasattr(self, '_store'):
return
self.close()
def __iter__(self):
for k in self.keys():
yield k
def __len__(self):
return len(self.keys())
def __repr__(self):
return repr(dict(self.iteritems()))
def close(self):
'''Closes internal store and clears object references.'''
try:
self._store.close()
except AttributeError:
pass
self._store = None
def clear(self):
'''Removes all keys and values from a store.'''
for key in self.keys():
del self[key]
def items(self):
'''Returns a list with all key/value pairs in the store.'''
return list(self.iteritems())
def iteritems(self):
'''Lazily returns all key/value pairs in a store.'''
for k in self:
yield (k, self[k])
def iterkeys(self):
'''Lazy returns all keys in a store.'''
return self.__iter__()
def itervalues(self):
'''Lazily returns all values in a store.'''
for _, v in self.iteritems():
yield v
def keys(self):
'''Returns a list with all keys in a store.'''
raise NotImplementedError()
def pop(self, key, *args):
'''
Removes and returns a value from a store.
@param args Default to return if key not present.
'''
if len(args) > 1:
raise TypeError('pop expected at most 2 arguments, got ' + repr(
1 + len(args))
)
try:
value = self[key]
# Return default if key not in store
except KeyError:
if args:
return args[0]
del self[key]
return value
def popitem(self):
'''Removes and returns a key, value pair from a store.'''
try:
k, v = self.iteritems().next()
except StopIteration:
raise KeyError('Store is empty.')
del self[k]
return (k, v)
def setdefault(self, key, default=None):
'''
Returns the value corresponding to an existing key or sets the
to key to the default and returns the default.
@param default Default value (default: None)
'''
try:
return self[key]
except KeyError:
self[key] = default
return default
def update(self, other=None, **kw):
'''
Adds to or overwrites the values in this store with values from
another store.
other Another store
kw Additional keys and values to store
'''
if other is None:
pass
elif hasattr(other, 'iteritems'):
for k, v in other.iteritems():
self[k] = v
elif hasattr(other, 'keys'):
for k in other.keys():
self[k] = other[k]
else:
for k, v in other:
self[k] = v
if kw:
self.update(kw)
def values(self):
'''Returns a list with all values in a store.'''
return list(v for _, v in self.iteritems())
class Shove(BaseStore):
'''Common object frontend class.'''
def __init__(self, store='simple://', cache='simple://', **kw):
super(Shove, self).__init__(store, **kw)
# Load store
self._store = getbackend(store, stores, **kw)
# Load cache
self._cache = getbackend(cache, caches, **kw)
# Buffer for lazy writing and setting for syncing frequency
self._buffer, self._sync = dict(), kw.get('sync', 2)
def __getitem__(self, key):
'''Gets a item from shove.'''
try:
return self._cache[key]
except KeyError:
# Synchronize cache and store
self.sync()
value = self._store[key]
self._cache[key] = value
return value
def __setitem__(self, key, value):
'''Sets an item in shove.'''
self._cache[key] = self._buffer[key] = value
# When the buffer reaches self._limit, writes the buffer to the store
if len(self._buffer) >= self._sync:
self.sync()
def __delitem__(self, key):
'''Deletes an item from shove.'''
try:
del self._cache[key]
except KeyError:
pass
self.sync()
del self._store[key]
def keys(self):
'''Returns a list of keys in shove.'''
self.sync()
return self._store.keys()
def sync(self):
'''Writes buffer to store.'''
for k, v in self._buffer.iteritems():
self._store[k] = v
self._buffer.clear()
def close(self):
'''Finalizes and closes shove.'''
# If close has been called, pass
if self._store is not None:
try:
self.sync()
except AttributeError:
pass
self._store.close()
self._store = self._cache = self._buffer = None
class FileBase(Base):
'''Base class for file based storage.'''
def __init__(self, engine, **kw):
super(FileBase, self).__init__(engine, **kw)
if engine.startswith('file://'):
engine = urllib.url2pathname(engine.split('://')[1])
self._dir = engine
# Create directory
if not os.path.exists(self._dir):
self._createdir()
def __getitem__(self, key):
# (per Larry Meyn)
try:
item = open(self._key_to_file(key), 'rb')
data = item.read()
item.close()
return self.loads(data)
except:
raise KeyError(key)
def __setitem__(self, key, value):
# (per Larry Meyn)
try:
item = open(self._key_to_file(key), 'wb')
item.write(self.dumps(value))
item.close()
except (IOError, OSError):
raise KeyError(key)
def __delitem__(self, key):
try:
os.remove(self._key_to_file(key))
except (IOError, OSError):
raise KeyError(key)
def __contains__(self, key):
return os.path.exists(self._key_to_file(key))
def __len__(self):
return len(os.listdir(self._dir))
def _createdir(self):
'''Creates the store directory.'''
try:
os.makedirs(self._dir)
except OSError:
raise EnvironmentError(
'Cache directory "%s" does not exist and ' \
'could not be created' % self._dir
)
def _key_to_file(self, key):
'''Gives the filesystem path for a key.'''
return os.path.join(self._dir, urllib.quote_plus(key))
def keys(self):
'''Returns a list of keys in the store.'''
return [urllib.unquote_plus(name) for name in os.listdir(self._dir)]
class SimpleBase(Base):
'''Single-process in-memory store base class.'''
def __init__(self, engine, **kw):
super(SimpleBase, self).__init__(engine, **kw)
self._store = dict()
def __getitem__(self, key):
try:
return self._store[key]
except:
raise KeyError(key)
def __setitem__(self, key, value):
self._store[key] = value
def __delitem__(self, key):
try:
del self._store[key]
except:
raise KeyError(key)
def __len__(self):
return len(self._store)
def keys(self):
'''Returns a list of keys in the store.'''
return self._store.keys()
class LRUBase(SimpleBase):
def __init__(self, engine, **kw):
super(LRUBase, self).__init__(engine, **kw)
self._max_entries = kw.get('max_entries', 300)
self._hits = 0
self._misses = 0
self._queue = deque()
self._refcount = dict()
def __getitem__(self, key):
try:
value = super(LRUBase, self).__getitem__(key)
self._hits += 1
except KeyError:
self._misses += 1
raise
self._housekeep(key)
return value
def __setitem__(self, key, value):
super(LRUBase, self).__setitem__(key, value)
self._housekeep(key)
if len(self._store) > self._max_entries:
while len(self._store) > self._max_entries:
k = self._queue.popleft()
self._refcount[k] -= 1
if not self._refcount[k]:
super(LRUBase, self).__delitem__(k)
del self._refcount[k]
def _housekeep(self, key):
self._queue.append(key)
self._refcount[key] = self._refcount.get(key, 0) + 1
if len(self._queue) > self._max_entries * 4:
self._purge_queue()
def _purge_queue(self):
for i in [None] * len(self._queue):
k = self._queue.popleft()
if self._refcount[k] == 1:
self._queue.append(k)
else:
self._refcount[k] -= 1
class DbBase(Base):
'''Database common base class.'''
def __init__(self, engine, **kw):
super(DbBase, self).__init__(engine, **kw)
def __delitem__(self, key):
self._store.delete(self._store.c.key == key).execute()
def __len__(self):
return self._store.count().execute().fetchone()[0]
__all__ = ['Shove']