root/branches/crazycache/dejavu/storage/storememcached.py

import md5
import memcache
import re
import sys

try:
    set
except NameError:
    from sets import Set as set

import dejavu
from dejavu import errors, logflags, storage
from geniusql import logic


def bytecode_regex(bits):
    """Make a regular expression out of the given mixed bytecode bits.
    
    If any bit is an integer, it will be replaced with re.escape(chr(bit)).
    Any bits which are already strings will be added as-is.
    """
    s = []
    for bit in bits:
        if not isinstance(bit, basestring):
            bit = re.escape(chr(bit))
        s.append(bit)
    return "".join(s)

simple_compare = bytecode_regex([124, 0, 0, 105, ".", ".", 100, ".", ".", 106, 2, 0])
simple_and = bytecode_regex([111, ".", ".", 1])
indexable_regex = re.compile("^(%s(%s)?)+S$" % (simple_compare, simple_and))


class MemcachedStorageManager(storage.StorageManager):
    """A Storage Manager which keeps all data in memcached.
    
    memcached is a high-performance, distributed memory object caching
    system, generic in nature, but intended for use in speeding up
    dynamic web applications by alleviating database load.
    
    See http://www.danga.com/memcached/
    and ftp://ftp.tummy.com/pub/python-memcached/
    
    IMPORTANT: data stuck into memcached is not guaranteed to be stable.
    It may disappear at any time, according to an internal LRU algorithm.
    In particular, you should be aware that the LRU algorithm is itself
    partitioned by object size (into "slabs"), so that a newer object
    may be removed before an older one if they are of significantly
    different sizes.
    
    Options:
        memcached.servers: a list of strings of the form 'IP-address:port'.
            These will be passed directly into the memcache.Client instance.
        
        memcached.global_index: if True (the default), this store will
            maintain a index over the identifiers of all stored objects in
            memcached itself. This is the 'safe' choice, and necessary if
            your only store is memcached. However, if you run this store as
            an ObjectCache.cache, you should turn this off, allowing
            ObjectCache.nextstore to maintain the primary indexes--this
            allows the cache to run orders of magnitude faster.
        
        memcached.index_time: the timeout, in seconds, for any cached indexes.
            Default is 300 seconds.
    """
    
    def __init__(self, allOptions={}):
        storage.StorageManager.__init__(self, allOptions)
        
        self.name = allOptions['name']
        self.global_index = allOptions.pop("memcached.global_index", True)
        self.index_time = allOptions.pop("memcached.index_time", 5 * 60)
        self.primary_keys = {}
        self.indexsets = {}
        self.index_stride = 50
        
        cache_opts = dict([(k[10:], v) for k, v in allOptions.iteritems()
                           if k.startswith("memcached.")])
        self.client = memcache.Client(**cache_opts)
    
    def hash(self, object):
        """Return a consistent hash for object (for use in a memcached key)."""
        # TODO: can we add overflow support for collisions?
        return md5.new(repr(object)).hexdigest()
    
    def _unit_key(self, unit):
        """Return (ident, memcached key) for the given unit."""
        cls = unit.__class__
        ident = tuple([getattr(unit, name) for name in self.primary_keys[cls]])
        key = "%s:%s:%s" % (self.name, cls.__name__, self.hash(ident))
        return key
    
    def unit(self, cls, **kwargs):
        """A single Unit which matches the given kwargs, else None.
        
        The first Unit matching the kwargs is returned; if no Units match,
        None is returned.
        """
        keyset = set(kwargs.keys())
        
        # Try to retrieve a matching unit using its primary_keys.
        # This will skip grabbing any indices (a HUGE optimization).
        pk = self.primary_keys[cls]
        if keyset >= set(pk):
            return self._unit_by_primary_key(cls, pk, kwargs)
        
        # Try to retrieve a matching unit using an index.
        # If self.global_index is True, the last one should
        # be an index with propnames == []. See self.register.
        indexset = self.indexsets[cls]
        for index in indexset:
            if keyset >= set(index):
                unit = indexset.unit(index, kwargs)
                if unit is not None:
                    if self.logflags & logflags.RECALL:
                        self.log(logflags.RECALL.message(cls, ('HIT', kwargs)))
                    return unit
        
        # Return None since we have no more access paths.
        if self.logflags & logflags.RECALL:
            self.log(logflags.RECALL.message(cls, ('DEFER', kwargs)))
        return None
    
    def xrecall(self, classes, expr=None, order=None, limit=None, offset=None):
        """Yield units of the given cls which match the given expr."""
        if isinstance(classes, dejavu.UnitJoin):
            for units in self._xmultirecall(classes, expr, order=order,
                                            limit=limit, offset=offset):
                yield units
            return
        
        cls = classes
        indexset = self.indexsets[cls]
        
        if not isinstance(expr, logic.Expression):
            expr = logic.Expression(expr)
        if self.logflags & logflags.RECALL:
            self.log(logflags.RECALL.message(cls, expr))
        
        if limit == 0:
            return
        
        if offset and not order:
            raise TypeError("Order argument expected when offset is provided.")
        
        filters = self.extract_filters(expr)
        
        # Try to retrieve a single matching unit using its primary_keys.
        # This will skip grabbing any indices (a HUGE optimization).
        pk = self.primary_keys[cls]
        if set(filters.keys()) >= set(pk):
            yield self._unit_by_primary_key(cls, pk, filters)
            return
        
        # Try to retrieve matching units using an index.
        # If self.global_index is True, the last one should
        # be an index with propnames == []. See self.register.
        for index in indexset:
            if set(filters.keys()) >= set(index):
                data = indexset.xrecall(index, filters)
                data = self._xrecall_inner(data, expr)
                for unit in self._paginate(data, order, limit, offset, single=True):
                    yield unit
                return
    
    def _xrecall_inner(self, units, expr=None):
        """Private helper for self.xrecall."""
        for unit in units:
            if expr is None or expr(unit):
                # Must yield a sequence for use in _paginate.
                yield (unit,)
    
    def save(self, unit, forceSave=False):
        """Store the unit."""
        if self.logflags & logflags.SAVE:
            self.log(logflags.SAVE.message(unit, forceSave))
        
        if forceSave or unit.dirty():
            # Cleanse first because pickle state
            # includes _initial_property_hash.
            unit.cleanse()
            self.client.set(self._unit_key(unit), unit)
            self.indexsets[unit.__class__].add(unit)
    
    def destroy(self, unit):
        """Delete the unit."""
        if self.logflags & logflags.DESTROY:
            self.log(logflags.DESTROY.message(unit))
        
        self.client.delete(self._unit_key(unit))
        self.indexsets[unit.__class__].discard(unit)
    
    def reserve(self, unit):
        """Reserve storage space for the Unit."""
        if unit.identifiers:
            cls = unit.__class__
            indexset = self.indexsets[cls]
            
            if not unit.sequencer.valid_id(unit.identity()):
                if () in indexset:
                    # Try to generate an identifier by looking
                    # up all units in the global index.
                    index = indexset.get({}) or []
                    ids = [u.identity()
                           for u in indexset.scan(index).itervalues()]
                    unit.sequencer.assign(unit, ids)
                else:
                    raise NotImplementedError(
                        "Unindexed memcache cannot generate identifiers.")
            
            unit.cleanse()
            
            # Add the unit to the cache.
            try:
                self.client.add(self._unit_key(unit), unit)
            except IOError, exc:
                if exc.args[0] == 'NOT_STORED':
                    pass
                raise
            
            # Add the unit to all indices.
            indexset.add(unit)
        else:
            # This class has no identifiers, so skip reserve and wait for save.
            pass
        
        # Usually we log ASAP, but here we log after
        # the unit has had a chance to get an auto ID.
        if self.logflags & logflags.RESERVE:
            self.log(logflags.RESERVE.message(unit))
    
    def shutdown(self, conflicts='error'):
        """Shut down all connections to internal storage.
        
        conflicts: see errors.conflict.
        """
        self.client.disconnect_all()
    
    def create_database(self, conflicts='error'):
        """Create internal structures for the entire database.
        
        conflicts: see errors.conflict.
        """
        pass
    
    def drop_database(self, conflicts='error'):
        """Destroy internal structures for the entire database.
        
        conflicts: see errors.conflict.
        """
        for cls in self.classes:
            self.flush(cls)
    
    def create_storage(self, cls, conflicts='error'):
        """Create internal structures for the given class.
        
        conflicts: see errors.conflict.
        """
        if self.logflags & logflags.DDL:
            self.log(logflags.DDL.message("create storage %s" % cls))
        
        indexset = self.indexsets[cls]
        if () in indexset:
            try:
                self.client.add(indexset.key({}), [])
            except IOError, exc:
                if exc.args[0] == 'NOT STORED':
                    errors.conflict(conflicts, "Class %r already has storage."
                                    % cls)
                else:
                    raise
    
    def has_storage(self, cls):
        """If storage structures exist for the given class, return True."""
        return True
    
    def drop_storage(self, cls, conflicts='error'):
        """Destroy internal structures for the given class.
        
        conflicts: see errors.conflict.
        """
        if self.logflags & logflags.DDL:
            self.log(logflags.DDL.message("drop storage %s" % cls))
        self.flush(cls)
    
    def add_property(self, cls, name, conflicts='error'):
        """Add internal structures for the given property.
        
        conflicts: see errors.conflict.
        """
        clsname = cls.__name__
        if self.logflags & logflags.DDL:
            self.log(logflags.DDL.message("add property %s %s" %
                                          (clsname, name)))
        
        indexset = self.indexsets[cls]
        if () in indexset:
            # TODO: recalculate if primary_keys changed
            ci = self.client.get(indexset.key({})) or []
            for id in ci:
                key = "%s:%s:%s" % (self.name, clsname, self.hash(id))
                unit = self.client.get(key)
                if unit is not None:
                    unit._properties[name] = None
                    unit.cleanse()
                    self.client.set(key, unit)
    
    def has_property(self, cls, name):
        """If storage structures exist for the given property, return True."""
        indexset = self.indexsets[cls]
        if () in indexset:
            clsname = cls.__name__
            ci = self.client.get(indexset.key({}))
            
            if not ci:
                # We don't have any items, so there's nothing to
                # declare as 'unprepared'.
                return True
            
            for id in ci:
                key = "%s:%s:%s" % (self.name, clsname, self.hash(id))
                unit = self.client.get(key)
                if unit is not None:
                    return name in unit._properties
        
        return True
    
    def drop_property(self, cls, name, conflicts='error'):
        """Destroy internal structures for the given property.
        
        conflicts: see errors.conflict.
        """
        clsname = cls.__name__
        if self.logflags & logflags.DDL:
            self.log(logflags.DDL.message("drop property %s %s" %
                                          (clsname, name)))
        
        indexset = self.indexsets[cls]
        if () in indexset:
            ci = self.client.get(indexset.key({})) or []
            for id in ci:
                key = "%s:%s:%s" % (self.name, clsname, self.hash(id))
                unit = self.client.get(key)
                if unit is not None:
                    del unit._properties[name]
                    unit.cleanse()
                    self.client.set(key, unit)
    
    def rename_property(self, cls, oldname, newname, conflicts='error'):
        """Rename internal structures for the given property.
        
        conflicts: see errors.conflict.
        """
        clsname = cls.__name__
        if self.logflags & logflags.DDL:
            self.log(logflags.DDL.message("rename property %s from %s to %s"
                                          % (cls, oldname, newname)))
        
        indexset = self.indexsets[cls]
        if () in indexset:
            ci = self.client.get(indexset.key({})) or []
            for id in ci:
                key = "%s:%s:%s" % (self.name, clsname, self.hash(id))
                unit = self.client.get(key)
                if unit is not None:
                    unit._properties[newname] = unit._properties[oldname]
                    del unit._properties[oldname]
                    unit.cleanse()
                    self.client.set(key, unit)
    
    
    #                   Extra methods for use as a cache                   #
    
    def cachelen(self, cls):
        indexset = self.indexsets[cls]
        if () in indexset:
            return len(self.client.get(indexset.key({})))
        else:
            return 0
    
    def cached_units(self, cls):
        units = []
        indexset = self.indexsets[cls]
        if () in indexset:
            for key in self.client.get(indexset.key({})):
                unit = self.client.get(key)
                if unit is not None:
                    unit.cleanse()
                    units.append(unit)
        return units
    
    def flush(self, cls):
        """Dump all objects of the given class."""
        clsname = cls.__name__
        
        indexset = self.indexsets[cls]
        if () in indexset:
            gi_key = indexset.key({})
            # Delete all units in the global index.
            for id in self.client.get(gi_key) or []:
                key = "%s:%s:%s" % (self.name, clsname, self.hash(id))
                self.client.delete(key)
            
            # Delete the global index.
            self.client.delete(gi_key)
        # TODO:
        # else:
        #     self.increment_generation(cls)
    
    def register(self, cls):
        """Assert that Units of class 'cls' will be handled."""
        # Set a default primary key for the class. Consumers are free to
        # change this if another unique property is looked up more often.
        self.primary_keys[cls] = tuple(cls.identifiers or cls.properties)
        
        # Add indices based on the .index attribute of each UnitProperty.
        self.indexsets[cls] = i = IndexSet(self, cls)
        for propname in cls.properties:
            prop = getattr(cls, propname)
            if prop.index:
                # There's usually no need for an index on the primary key;
                # we can just fetch each one directly by cache key.
                # Callers are free to add one in explicitly if needed,
                # for example, if wanting to retrive all units without
                # any filtering criteria.
                if propname not in cls.identifiers:
                    i.add_index(propname)
        
        # Add an index with no propnames. This is a special
        # sentinel value for the global index that keeps us DRY.
        if self.global_index:
            i.add_index()
        
        storage.StorageManager.register(self, cls)
    
    def _unit_by_primary_key(self, cls, keys, filters):
        """Return a unit (or None) by primary keys which matches the filters dict.
        
        The filters argument must contain an entry for each key in the
        given list of keys, although it may and often should contain
        additional entries.
        """
        ident = tuple([filters[k] for k in keys])
        key = "%s:%s:%s" % (self.name, cls.__name__, self.hash(ident))
        unit = self.client.get(key)
        if unit is not None:
            matching = True
            if set(filters.keys()) > set(keys):
                # We retrieved the Unit using a subset of the filters.
                # Filter in full now.
                for k, v in filters.iteritems():
                    if getattr(unit, k) != v:
                        matching =