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

"""Base classes and tools for writing database Storage Managers.

DATA TYPES
==========
Database Storage Manager modules are mostly adapters to support round-trip
data coercion:

Unit type -> [SQL repr ->] DB -> incoming Python value -> Unit type

Since Dejavu relies on external database servers for its persistence,
Python datatypes must be converted to column types in the DB. When writing
a StorageManager, you should make sure that your type conversions can handle
at least the following limitations: If possible, implement the type with no
limits. Also, follow UnitProperty.hints['bytes'] where possible. A value
of zero for hints['bytes'] implies no limit. If no value is given, try to
assume no limit, although you may choose whatever default size you wish
(255 is common for strings).

ENCODING ISSUES
===============
All SQL sent to the database must be strings, not unicode. You can set the
encoding of the Adapters (I may add a more centralized encoding context in
the future). We must use encoded strings so that we can mix encodings
within the same string; for example, we might have a DB which understands
utf8, but a pickle value which will be encoded in raw-unicode-escape inline
with that. All values, therefore, must be coerced before we try to join
them into an SQL statement string.

"""

import threading
import warnings


import geniusql
from geniusql import logic, logicfuncs

import dejavu
from dejavu import storage, logflags, xray
from dejavu.errors import StorageWarning, MappingError, conflict


# --------------------------- Storage Manager --------------------------- #


class StorageManagerDB(storage.StorageManager):
    """StoreManager base class to save and retrieve Units using a DB."""
    
    databaseclass = geniusql.Database
    
    def __init__(self, allOptions={}):
        storage.StorageManager.__init__(self, allOptions)
        self.reserve_lock = threading.Lock()
        
        # Config Overrides
        def get_option(name):
            item = allOptions.get(name)
            if isinstance(item, basestring):
                item = xray.classes(item)
            return item
        
        dbclass = get_option('Database Class')
        if dbclass:
            self.databaseclass = dbclass
        
        allOptions = dict([(str(k), v) for k, v in allOptions.iteritems()])
        
        self.db = self.databaseclass(**allOptions)
        self.schema = self.db.schema()
        
        if 'Prefix' in allOptions:
            self.schema.prefix = allOptions['Prefix']
        
        def logger(msg):
            if self.logflags & logflags.SQL:
                self.log(logflags.SQL.message(msg))
        self.db.log = logger
    
    def version(self):
        return self.db.version()
    
    def shutdown(self, conflicts='error'):
        """Shut down all connections to internal storage.
        
        conflicts: see errors.conflict.
        """
        self.db.connections.shutdown()
    
    def xrecall(self, classes, expr=None, order=None, limit=None, offset=None):
        """Yield a sequence of Unit instances which satisfy the expression."""
        if isinstance(classes, dejavu.UnitJoin):
            for unitrow in self._xmultirecall(classes, expr, order=order,
                                              limit=limit, offset=offset):
                yield unitrow
            return
        else:
            cls = classes
        
        if self.logflags & logflags.RECALL:
            self.log(logflags.RECALL.message(cls, expr))
        
        clsname = cls.__name__
        
        # Put the identifier properties first, in case other fields
        # depend upon them.
        idnames = list(cls.identifiers)
        attrs = idnames + [x for x in cls.properties if x not in idnames]
        coercers = [getattr(cls, key).coerce for key in attrs]
        
        data = self.select((cls, attrs, expr), order=order,
                           limit=limit, offset=offset)
        for row in data:
            unit = cls.__new__(cls)
            unit._zombie = True
            unit.__init__()
            
            for key, value, propcoerce in zip(attrs, row, coercers):
                try:
                    if propcoerce:
                        value = propcoerce(unit, value)
                    unit._properties[key] = value
                except UnicodeDecodeError, x:
                    x.reason += " [%r: %r]" % (key, value)
                    raise
                except Exception, x:
                    x.args += (key, value)
                    raise
            
            # If our SQL is imperfect, don't yield it to the
            # caller unless it passes expr(unit).
            if expr and data.statement.imperfect:
                if not expr(unit):
                    continue
            
            unit.cleanse()
            yield unit
    
    def reserve(self, unit):
        """Reserve a persistent slot for unit."""
        self.reserve_lock.acquire()
        try:
            # First, see if our db subclass has a handler that
            # uses the DB to generate the appropriate identifier(s).
            seqclass = unit.sequencer.__class__.__name__
            seq_handler = getattr(self, "_seq_%s" % seqclass, None)
            if seq_handler:
                seq_handler(unit)
            else:
                self._manual_reserve(unit)
            unit.cleanse()
        finally:
            self.reserve_lock.release()
        
        # 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 _seq_UnitSequencerInteger(self, unit):
        """Reserve a unit (using the table's autoincrement fields)."""
        cls = unit.__class__
        
        # Grab the new ID. This is threadsafe because reserve has a mutex.
        newids = self.schema[cls.__name__].insert(**unit._properties)
        for k, v in newids.iteritems():
            setattr(unit, k, v)
    
    def _manual_reserve(self, unit):
        """Use when the DB cannot automatically generate an identifier.
        The identifiers will be supplied by UnitSequencer.assign().
        """
        cls = unit.__class__
        t = self.schema[cls.__name__]
        if not unit.sequencer.valid_id(unit.identity()):
            # Examine all existing IDs and grant the "next" one.
            data = list(self.db.select((t, cls.identifiers)))
            cls.sequencer.assign(unit, data)
        t.insert(**unit._properties)
    
    def save(self, unit, forceSave=False):
        """Update storage from unit's data (if unit.dirty())."""
        if self.logflags & logflags.SAVE:
            self.log(logflags.SAVE.message(unit, forceSave))
        
        if forceSave or unit.dirty():
            self.schema[unit.__class__.__name__].save(**unit._properties)
            unit.cleanse()
    
    def destroy(self, unit):
        """Delete the unit."""
        if self.logflags & logflags.DESTROY:
            self.log(logflags.DESTROY.message(unit))
        
        table = self.schema[unit.__class__.__name__]
        table.delete(**unit._properties)
    
    
    #                                Views                                #
    
    def tablejoin(self, join):
        """Return a geniusql Join tree for the given UnitJoin."""
        t1, t2 = join.class1, join.class2
        
        if isinstance(t1, dejavu.UnitJoin):
            wt1 = self.tablejoin(t1)
        else:
            wt1 = self.schema[t1.__name__]
        
        if isinstance(t2, dejavu.UnitJoin):
            wt2 = self.tablejoin(t2)
        else:
            wt2 = self.schema[t2.__name__]
        
        uj = geniusql.Join(wt1, wt2, join.leftbiased)
        # if the original UnitJoin had a custom association path,
        # copy it to the new Join instance
        uj.path = join.path
        return uj
    
    def _geniusql_query(self, query):
        """Return a Geniusql Query object for the given Dejavu Query."""
        rel = query.relation
        if isinstance(rel, dejavu.UnitJoin):
            rel = self.tablejoin(rel)
        elif rel is None:
            # This is a Geniusql-ism: send the schema when we have no FROM.
            rel = self.schema
        else:
            rel = self.schema[rel.__name__]
        return geniusql.Query(rel, query.attributes, query.restriction)
    
    def _geniusql_statement(self, statement):
        """Return a Geniusql Statement object for the given Dejavu Statement."""
        return geniusql.Statement(
            self._geniusql_query(statement.query),
            order=statement.order, limit=statement.limit,
            offset=statement.offset, distinct=statement.distinct)
    
    def select(self, query, order=None, limit=None, offset=None, distinct=None):
        """Return a geniusql Dataset for the given Query object."""
        if not isinstance(query, dejavu.Query):
            query = dejavu.Query(*query)
        
        return self.db.select(self._geniusql_query(query),
                              order=order, limit=limit, offset=offset,
                              distinct=distinct, strict=False)
    
    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)
        
        self.db.insert_into(name, self._geniusql_query(query),
                            distinct=distinct)
        return Modeler(self.schema).make_class(name)
    
    def make_class(self, name):
        """Return a (new) Unit class for the given storage name."""
        return Modeler(self.schema).make_class(name)
    
    def xview(self, query, order=None, limit=None, offset=None, distinct=False):
        """Yield value tuples for the given query."""
        if not isinstance(query, dejavu.Query):
            query = dejavu.Query(*query)
        
        if self.logflags & logflags.VIEW:
            self.log(logflags.VIEW.message(query, distinct))
        
        data = self.select(query, order=order, limit=limit, offset=offset,
                           distinct=distinct)
        if data.statement.imperfect:
            # ^%$#@! There's no way to handle imperfect queries without
            # creating all involved Units, which defeats the performance
            # benefits of view.
            clsname = self.__class__.__name__
            warnings.warn("The requested query cannot produce perfect SQL "
                          "with a %s datasource. It may take an absurdly "
                          "long time to run, since each unit must be fully-"
                          "formed. %s" % (clsname, query), StorageWarning)
            for row in storage.StorageManager.xview(self, query, order=order,
                                                    limit=limit, offset=offset,
                                                    distinct=distinct):
                yield row
        else:
            # Use tuples for hashability
            for row in data:
                yield tuple(row)
    
    def count(self, cls, expr=None):
        """Number of Units of the given cls which match the given expr."""
        if cls.identifiers:
            uniq = cls.identifiers
        else:
            uniq = cls._properties.keys()
        # TODO: handle multiple args to count()
        counter = lambda x: [logicfuncs.count(getattr(x, uniq[0]))]
        
        query = dejavu.Query(cls, counter, expr)
        
        if self.logflags & logflags.VIEW:
            self.log(logflags.VIEW.message(query, False))
        
        data = self.select(query)
        if data.statement.imperfect:
            # ^%$#@! There's no way to handle imperfect queries without
            # creating all involved Units, which defeats the performance
            # benefits of view.
            clsname = self.__class__.__name__
            warnings.warn("The requested query cannot produce perfect SQL "
                          "with a %s datasource. It may take an absurdly "
                          "long time to run, since each unit must be fully-"
                          "formed. %s" % (clsname, query), StorageWarning)
            return storage.StorageManager.count(self, cls, expr)
        else:
            return data.scalar()
    
    def _xmultirecall(self, classes, expr=None, order=None, limit=None, offset=None):
        """Yield Unit instance sets which satisfy the expression."""
        if self.logflags & logflags.RECALL:
            self.log(logflags.RECALL.message(classes, expr))
        
        # Gather attribute list.
        allattrs = []
        props = []
        for cls in classes:
            attrs = []
            for key in cls.properties:
                attrs.append(key)
                props.append((cls, key, getattr(cls, key).coerce))
            allattrs.append(attrs)
        
        data = self.select((classes, allattrs, expr), order=order,
                           limit=limit, offset=offset)
        for row in data:
            # TODO: This is broken; won't work if same cls appears twice.
            units = {}
            for i, (cls, key, propcoerce) in enumerate(props):
                if cls in units:
                    unit = units[cls]
                else:
                    unit = cls.__new__(cls)
                    unit._zombie = True
                    unit.__init__()
                    units[cls] = unit
                
                value = row[i]
                try:
                    if propcoerce:
                        value = propcoerce(unit, value)
                    unit._properties[key] = value
                except Exception, x:
                    x.args += (cls, key)
                    raise
            
            unitset = []
            for cls in classes:
                unit = units[cls]
                unit.cleanse()
                unitset.append(unit)
            
            # If our SQL is imperfect, don't yield units to the
            # caller unless they pass expr(unit).
            acceptable = True
            if expr and data.statement.imperfect:
                acceptable = expr(*unitset)
            if acceptable:
                yield unitset
    
    #                               Schemas                               #
    
    def create_database(self, conflicts='error'):
        """Create internal structures for the entire database.
        
        conflicts: see errors.conflict.
        """
        if self.logflags & logflags.DDL:
            self.log(logflags.DDL.message("create database"))
        
        try:
            self.db.create()
            self.schema.create()
        except geniusql.errors.MappingError, x:
            conflict(conflicts, str(x))
    
    def drop_database(self, conflicts='error'):
        """Destroy internal structures for the entire database.
        
        conflicts: see errors.conflict.
        """
        if self.logflags & logflags.DDL:
            self.log(logflags.DDL.message("drop database"))
        
        try:
            self.schema.drop()
        except geniusql.errors.MappingError, x:
            conflict(conflicts, str(x))
        
        try:
            self.db.drop()
        except geniusql.errors.MappingError, x:
            conflict(conflicts, str(x))
    
    def _make_table(self, cls):
        """Create and return a Table or View object for the given class."""
        if hasattr(cls, "view_statement"):
            gs = self._geniusql_statement(cls.view_statement)
            t = self.schema.view(cls.__name__, gs)
            fields = []
            for key in cls.properties:
                t[key] = self._make_column(cls, key)
        else:
            t = self.schema.table(cls.__name__)
            indices = cls.indices()
            fields = []
            for key in cls.properties:
                t[key] = col = self._make_column(cls, key)
                if key in indices:
                    t.add_index(key)
                if col.autoincrement and col.sequence_name is None:
                    # Not every database needs/uses sequence_name
                    col.sequence_name = self.schema.sequence_name(t.name, key)
        
        # Copy associations to table.references.
        for k, v in cls._associations.iteritems():
            t.references[k] = (v.nearKey, v.farClass.__name__, v.farKey)
        
        return t
    
    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))
        
        try:
            # Attach to self.schema, which should call CREATE TABLE.
            self.schema[cls.__name__] = self._make_table(cls)
        except geniusql.errors.MappingError,