root/trunk/dejavu/storage/partitions.py

"""A StorageManager for Dejavu which mediates multiple stores."""

from geniusql import logic

import dejavu
from dejavu import errors, storage, logflags


class VerticalPartitioner(storage.StorageManager):
    """A mediator for multiple vertically-partitioned stores.
    
    Options:
        <Name> Store: Use the name to look up another store and attach it
            to self.stores.
        <Name> Units: Use the name to look up another store and direct
            the given Unit classes to that named store.
    """
    
    # A map from store names to StorageManager instances
    stores = {}
    
    # A map from Unit classes to StorageManager instances
    classmap = {}
    
    # A map from Unit class names to StorageManager instances
    classnames = {}
    
    def __init__(self, allOptions={}):
        storage.StorageManager.__init__(self, allOptions)
        self.stores = {}
        self.classmap = {}
        self.classnames = {}
        self.default_store = None
        
        opts = allOptions.items()
        opts.sort()
        for k, v in opts:
            if k.endswith(" Store"):
                self.stores[k.split(" ", 1)[0]] = v
            elif k.endswith(" Units"):
                sm = self.stores[k.split(" ", 1)[0]]
                units = [x.strip() for x in v.split(",") if x.strip()]
                if units:
                    for x in units:
                        self.classnames[x] = sm
                else:
                    self.default_store = sm
    
    def migrate_class(self, cls, new_store):
        """Copy all units of cls to new_store."""
        new_store.create_storage(cls)
        for unit in self.new_sandbox().xrecall(cls):
            new_store.reserve(unit)
            new_store.save(unit, forceSave=True)
    
    def migrate(self, new_store, old_store=None, copy_only=False):
        """Copy all units (of old_store) to new_store."""
        for cls in self.classes:
            store = self.storage(cls)
            if old_store is None or old_store is store:
                self.migrate_class(cls, new_store)
                if not copy_only:
                    self.classmap[cls] = new_store
    
    def add_store(self, name, store, units=None):
        """Register a StorageManager to be mediated.
        
        name: a key for the given store.
        store: a StorageManager class, an instance of that class,
            or the full importable dotted-package name of the class.
        units: a list of strings, naming the classes this store
            will handle. If None, the store will handle all classes
            which are not handled by any other store.
        """
        self.stores[name] = store = storage.resolve(store, options)
        if units:
            for x in units:
                self.classnames[x] = store
        else:
            self.default_store = store
        return store
    
    def remove_store(self, name):
        """Remove (unregister) the named store.
        
        All classes associated to the given store will be disassociated,
        and will then be free to associate with another store (usually
        the default store).
        """
        if name in self.stores:
            store = self.stores[name]
            
            # Disassociate all registered classes with this store.
            for c in self.classes:
                if self.classmap[c] is store:
                    self.classmap[c] = None
            
            if store is self.default_store:
                self.default_store = None
            
            del self.stores[name]
    
    def map_all(self, conflict_mode='error'):
        """Map all registered classes to internal storage structures.
        
        Although classes are mapped automatically the first time they
        are accessed (see self.storage), in production it is often more
        useful to map all classes at application startup. Call this
        method to do so (but register all classes first).
        
        This method is idempotent, but that doesn't mean cheap. Try not
        to call it very often (once at app startup is usually enough).
        
        conflict_mode: This argument determines what happens when there are
        discrepancies between the Dejavu model and the actual database.
            
            If 'error' (the default), MappingError is raised for the
            first issue and the sync process is aborted.
            
            If 'warn', then a warning is raised (instead of an error)
            for each issue, and the sync process is not aborted. This
            allows you to see all errors at once, without having to stop
            and fix each one and then execute the process again.
            
            If 'repair', then each issue will be resolved by changing
            the database to match the model.
        """
        storemap = {}
        
        # Don't use iteritems because self.storage may mutate the dict.
        for cls, store in self.classmap.items():
            if store is None:
                # Find storage directly to skip the map() call in self.storage.
                store = self.classnames.get(cls.__name__, self.default_store)
                self.classmap[cls] = store
            bucket = storemap.setdefault(store, [])
            bucket.append(cls)
        
        storemap = [(s.loadOrder, s, c) for s, c in storemap.iteritems()]
        storemap.sort()
        
        for order, store, classes in storemap:
            store.map(classes, conflict_mode)
    
    def shutdown(self):
        """Shutdown self and all its stores."""
        # Tell all stores to shut down.
        stores = [(v.shutdownOrder, v, k) for k, v in self.stores.iteritems()]
        stores.sort()
        for order, store, name in stores:
            store.shutdown()
    
    def version(self):
        """Return provider-specific version strings for each mediated store."""
        output = []
        for store in self.stores.itervalues():
            if store.version:
                output.append(store.version())
        return '\n\n'.join(output)
    
    
    # --------------------- Unit Class Registration --------------------- #
    
    def create_database(self):
        for s in self.stores.itervalues():
            s.create_database()
    
    def drop_database(self):
        for s in self.stores.itervalues():
            s.drop_database()
    
    def register(self, cls):
        """Assert that Units of class 'cls' will be handled."""
        storage.StorageManager.register(self, cls)
        
        # We must allow modules to register classes before any stores have
        # been added, but not overwrite a store which has already been found.
        if cls not in self.classmap:
            self.classmap[cls] = None
    
    def storage(self, cls, automap=True):
        """Return the StorageManager which handles Units of the given class.
        
        The results of this call will be cached for performance.
        
        If automap is True (the default), this method will call
        store.map(cls) if the class is not already cached.
        """
        store = self.classmap.get(cls)
        if store:
            return store
        
        store = self.classnames.get(cls.__name__, self.default_store)
        self.classmap[cls] = store
        if automap:
            store.map([cls])
        return store
    
    def create_storage(self, cls):
        """Create storage space for cls."""
        store = self.storage(cls, automap=False)
        store.create_storage(cls)
    
    def has_storage(self, cls):
        """If storage space for cls exists, return True (False otherwise)."""
        try:
            store = self.storage(cls)
        except errors.MappingError:
            return False
        return store.has_storage(cls)
    
    def drop_storage(self, cls):
        """Remove storage space for cls."""
        self.storage(cls).drop_storage(cls)
    
    def add_property(self, cls, name):
        """Add storage space for the named property of the given cls."""
        self.storage(cls).add_property(cls, name)
    
    def drop_property(self, cls, name):
        """Drop storage space for the named property of the given cls."""
        self.storage(cls).drop_property(cls, name)
    
    def rename_property(self, cls, oldname, newname):
        """Rename storage space for the property of the given cls."""
        self.storage(cls).rename_property(cls, oldname, newname)
    
    
    # ------------------------------- DML ------------------------------- #
    
    def reserve(self, unit):
        """Reserve storage space for the Unit."""
        self.storage(unit.__class__).reserve(unit)
    
    def save(self, unit, forceSave=False):
        """Store the unit's property values."""
        self.storage(unit.__class__).save(unit, forceSave)
    
    def destroy(self, unit):
        """Delete the unit."""
        self.storage(unit.__class__).destroy(unit)
    
    def xrecall(self, cls, expr=None, order=None, limit=None, offset=None):
        """Yield a sequence of Unit instances which satisfy the expression."""
        if isinstance(cls, dejavu.UnitJoin):
            for unitrow in self.multirecall(cls, expr):
                yield unitrow
        else:
            for unit in self.storage(cls).xrecall(cls, expr, order,
                                                  limit, offset):
                yield unit
    
    def xmultirecall(self, classes, expr=None, order=None, limit=None, offset=None):
        """Yield lists of units of the given classes which match expr.
        
        This does not yet handle multiple classes in disparate stores.
        """
        if expr is None:
            expr = logic.Expression(lambda *args: True)
        
        return self._single_store(classes).xmultirecall(classes, expr, order,
                                                        limit, offset)
    
    def _single_store(self, relation):
        """Return the store for the given relation (or raise ValueError)."""
        if hasattr(relation, "class1"):
            # This is a UnitJoin
            stores = [self.storage(cls) for cls in relation]
            store = stores[0]
            for s in stores:
                if s is not store:
                    raise ValueError(u"This operation does not support multiple"
                                     u" classes in disparate stores.")
        else:
            store = self.storage(relation)
        return store
    
    def xview(self, query, distinct=False):
        """Yield tuples of attribute values for the given query.
        
        Each yielded value will be a list of values, in the same order as
        the Query.attributes. This facilitates unpacking in iterative
        consumer code like:
        
        for id, name in store.view(Query(Invoice, ['ID', 'Name'], f)):
            print id, ": ", name
        
        This is generally much faster than recall, and should be preferred
        for performance-sensitive code.
        """
        if not isinstance(query, dejavu.Query):
            query = dejavu.Query(*query)
        
        if self.logflags & logflags.VIEW:
            self.log(logflags.VIEW.message(query, distinct))
        
        store = self._single_store(query.relation)
        for row in store.xview(query, distinct=distinct):
            yield row
    
    def insert_into(self, name, query, distinct=False):
        """INSERT matching data INTO a new class and return the class."""
        if not isinstance(query, dejavu.Query):
            query = dejavu.Query(*query)
        
        store = self._single_store(query.relation)
        return store.insert_into(name, query, distinct)
    
    #                        Transaction Management                        #
    
    def start(self, isolation=None):
        """Start a transaction."""
        for store in self.stores.itervalues():
            # By default, stores do not support transactions,
            # in which case 'start' will be None.
            if store.start:
                store.start(isolation)
    
    def commit(self):
        """Commit the current transaction.
        
        If errors occur during this process, they are not trapped here.
        You must either call rollback yourself (or fix the problem and
        try to commit again).
        """
        for store in self.stores.itervalues():
            if store.commit:
                store.commit()
    
    def rollback(self):
        """Roll back the current transaction."""
        for store in self.stores.itervalues():
            if store.rollback:
                store.rollback()