orm_tools

"""
Tools for persisting object graphs with Django ORM.

These classes make some assumptions about your datamodel:
Object graph does not have cyclic SQL insertion order dependencies.
Objects with pk=None are being inserted, all other pk values indicate an update.
"""

import threading

from django.db.models.base import Model
from django.db.models.query import QuerySet
from django.db.models.fields.related import ForeignKey

def is_managed(item):
    """Return True if item is managed by Django ORM."""
    
    return isinstance(item, Model)

def is_persistent(item):
    """Return True if an item has already been saved."""
    
    return item.pk is not None

class GraphError(Exception):
    pass

class Collection(object):
    """
    Allows iterables to be persisted in an object graph in an easier way.
    """
    
    @classmethod
    def set_property(cls, model, attr, parent_attr, set_attr=None):
        """
        Call this method with a model class as the 1st argument to
        create a property function which returns a Collection object.
        """
        
        private_attr = '_' + attr
    
        def _get(self):
            val = getattr(self, private_attr, None)
            if val is None:
                val = cls(self, parent_attr, set_attr=set_attr)
                setattr(self, private_attr, val)
            return val
    
        def _set(self, val):
            setattr(self, private_attr, val)
        setattr(model, attr, property(_get, _set))
    
        collection_attrs = getattr(model, '_collection_attrs', None)
        if collection_attrs is None:
            collection_attrs = []
            setattr(model, '_collection_attrs', collection_attrs)
        collection_attrs.append((attr, private_attr))
    
    def __init__(self, parent, parent_attr, set_attr=None, list_cls=None):
        self.parent = parent
        self.parent_attr = parent_attr
        self.set_attr = set_attr
        
        if list_cls is None:
            list_cls = list
        self.list_cls = list_cls
        
        # Holds list of items if parent is not persisted
        self._item_list = None
        
        # Holds set of items that have been accessed
        self._accessed_items = set()
        
    def __iter__(self):
        for item in self._items:
            self._accessed_items.add(item)
            yield item
            
    def __getitem__(self, k):
        if isinstance(k, slice):
            def _gen():
                items = self._items.__getitem__(k)
                for item in items:
                    self._accessed_items.add(item)
                    yield item
            return _gen()
        else:
            item = self._items.__getitem__(k)
            self._accessed_items.add(item)
            return item
    
    def __len__(self):
        items = self._items
        if isinstance(items, QuerySet):
            return items.count()
        else:
            return len(items)
    
    def _get_accessed_items(self):
        if is_persistent(self.parent) is True:
            # Return all items that have been accessed,
            # but only if they still belong in the collection.
            accessed_items = set()
            for item in self._accessed_items:
                parent = getattr(item, self.parent_attr, None)
                if parent is None:
                    continue
                
                if parent.pk != self.parent.pk:
                    continue
                
                accessed_items.add(item)
            return accessed_items
        else:
            return self._items
    accessed_items = property(_get_accessed_items)
    
    def _get_items(self):
        """Returns all items in the collection."""
        
        if is_persistent(self.parent) is True:
            return self.get_item_query()
        else:
            if self._item_list is None:
                self._item_list = self.list_cls()
            return self._item_list
    _items = property(_get_items)
    
    def add(self, item):
        """Add an item to the collection."""
        
        if is_managed(item) is False:
            # Bad things will happen
            raise GraphError('Cannot add unmanaged object.')
        
        setattr(item, self.parent_attr, self.parent)
        if is_persistent(self.parent) is True:
            item.save()
        else:
            found = False
            for existing_item in self._items:
                if existing_item is item:
                    found = True
                    break
            
            if found is False:
                self._items.append(item)
    
    def remove(self, item, delete=False):
        """Remove an item from the collection."""
        
        if delete is False:
            setattr(item, self.parent_attr, None)
            
        if is_persistent(self.parent) is True:
            if delete is True:
                item.delete()
            else:
                item.save()
        else:
            if item in self._items:
                self._items.remove(item)
    
    def clear(self, delete=False):
        """Remove all items from the collection."""
        
        to_remove = list(self._items)
        for item in to_remove:
            self.remove(item, delete=delete)
    
    def update(self, new):
        """[].update"""
        
        for item in new:
            self.add(item)
    
    def get_item_query(self):
        """Override this method to use a custom QuerySet."""
        
        return getattr(self.parent, self.set_attr).all()

class DependencyList(object):
    """Lists all parents for a child where parent == dependency."""
    
    __slots__ = ['obj', 'deps']
    
    def __init__(self, obj=None, deps=None):

        self.obj = obj
        self.deps = deps

class Dependency(object):
    """Signifies a dependency between two objects in the graph."""

    __slots__ = ['parent', 'field', 'level']

    def __init__(self, parent=None, field=None, level=None):

        self.parent = parent
        self.field = field
        self.level = level

class GraphSaver(object):
    """Provides methods to simplify the persistence of large object graphs."""

    def _add_dep(self, parent, child, field, deps, level):
        """Adds a dependency to the list."""

        dep_list = self._get_dep_list(child, deps)
        if dep_list is None:
            raise GraphError('Dependency list not found!')
        
        for dep in dep_list.deps:
            if dep.parent is parent:
                raise GraphError('Circular dependency detected.')

        dep_list.deps.append(Dependency(
            parent=parent,
            field=field,
            level=level))
        
    def _get_dep_list(self, obj, deps):
        """Returns a list that dependencies can be added to, or None"""
        
        dep_key = id(obj)
        return deps.get(dep_key, None)
    
    def _init_dep_list(self, obj, deps):
        """Returns a list that dependencies can be added to."""
        
        dep_key = id(obj)
        dep_list = deps.get(dep_key, None)
        dep_list = self._get_dep_list(obj, deps)
        if dep_list is None:
            dep_list = DependencyList(obj=obj, deps=[])
            deps[dep_key] = dep_list
        return dep_list

    def _build_deps(self, parent, deps, update=True, level=1):
        
        if self._get_dep_list(parent, deps) is not None:
            # This object has already had it's dependencies added!
            return

        if (is_persistent(parent) is False) or (update is True):
            # Makes sure the parent obj shows
            # up in the dependency list, so that
            # it gets saved!
            self._init_dep_list(parent, deps)

        for name in parent._meta.get_all_field_names():
            field_info = parent._meta.get_field_by_name(name)
            if field_info[2] is False:
                # This is a magic reverse reference to
                # some other related field. Ignore it.
                continue
            
            field = field_info[0]
            if isinstance(field, ForeignKey):
               child = getattr(parent, name)
               if child is not None:
                   # This is a dependency!
                   # Recursively add to deps.
                   self._build_deps(child, deps, update=update, level=level + 1)
                   self._add_dep(parent, child, field, deps, level)
                   
        # Save items in any collections
        self._build_collection_deps(parent, deps, update=update, level=level)
        
    def _build_collection_deps(self, parent, deps, update=True, level=1):
        """Add dependencies from collection objects."""
        
        collection_attrs = getattr(parent.__class__, '_collection_attrs', [])
        for collection_attr in collection_attrs:
            collection = getattr(parent, collection_attr[0], None)
            if collection is not None:
                for item in collection.accessed_items:
                    self._build_deps(item, deps, update=update, level=level)
    
    def _save_deps(self, deps):
        # Group dependencies by their level
        dep_levels = {}
        for dep_list in deps.itervalues():
            max_level = 0
            for dep in dep_list.deps:
                if dep.level > max_level:
                    max_level = dep.level
            
            level_list = dep_levels.get(max_level, None)
            if level_list is None:
                level_list = []
                dep_levels[max_level] = level_list
            level_list.append(dep_list)

        self._save_by_level(dep_levels)

    def _save_by_level(self, levels):
        # Save children by their dependency level
        level_keys = levels.keys()
        level_keys.sort()
        level_keys.reverse()
        for level in level_keys:
            for dep_list in levels[level]:
                self._save_dep_list(dep_list)

    def _save_dep_list(self, dep_list):
        # All children in the list should be the same,
        # so we only need to save the 1st one.
        child = dep_list.obj
        child.save()

        for dep in dep_list.deps:
            setattr(dep.parent, dep.field.name, child)

    def save_many(self, items, update=True):
        """Save multiple object graphs."""
        
        deps = {}
        for item in items:
            if is_managed(item) is False:
                raise GraphError('Cannot save unmanaged item.')
            self._build_deps(item, deps, level=1, update=update)
        self._save_deps(deps)

    def save(self, item, update=True):
        """Save an object graph."""
        
        return self.save_many((item,), update=update)

class Session(object):
    """
    Modifies QuerySets to return consistent object graphs.
    
    This will modify query behavior for all code executed
    during a 'with' statement.
    """
    
    _session_objs = threading.local()
    
    def __init__(self):
        self._entered = False
        self._teardown_on_exit = False
    
    def __enter__(self, *args, **kwargs):
        """Setup session system."""
        
        if self._entered is True:
            raise GraphError('Cannot re-enter session!')
        
        if self._is_setup() is False:
            self._setup()
            
        self._entered = True
            
    def __exit__(self, *args, **kwargs):
        if self._teardown_on_exit is True:
            self._teardown()
            
    def _is_setup(self):
        return getattr(self._session_objs, '_setup', False)
    
    def _setup(self):
        self._clear()
        self._patch_query_set()
        self._patch_model_base()
        self._session_objs._setup = True
        self._teardown_on_exit = True
        
    def _teardown(self):
        self._clear()
        self._session_objs._setup = False
    
    def _clear(self):
        """Clear all existing session objects."""
        
        self._session_objs.objs = {}
    
    def _patch_query_set(self):
        """Monkey patches QuerySet to return cached objects."""

        if getattr(QuerySet, '_session_enabled', False) is True:
            return

        f = QuerySet.iterator
        def _gen(q, *args, **kwargs):
            itr = f(q, *args, **kwargs)
            for item in itr:
                yield self.add(item)
        
        def _s_iterator(q, *args, **kwargs):
            if self._is_setup():
                return _gen(q, *args, **kwargs)
            else:
                return f(q, *args, **kwargs)
        QuerySet.iterator = _s_iterator
        QuerySet._session_enabled = True

    def _patch_model_base(self):
        """Monkey patches Model to added saved items to session."""
       
        if getattr(Model, '_session_enabled', False) is True:
            return
 
        f = Model.save
        def _s_save(item, *args, **kwargs):
            if self._is_setup():
                if is_persistent(item) is True:
                    # updating this object
                    if ('force_insert' not in kwargs) and ('force_update' not in kwargs):
                        kwargs['force_update'] = True
                    self.add_with_dup_check(item)
                    f(item, *args, **kwargs)
                else:
                    # inserting this object
                    if ('force_insert' not in kwargs) and ('force_update' not in kwargs):
                        kwargs['force_insert'] = True
                    f(item, *args, **kwargs)
                    self.add_with_dup_check(item)
            else:
                return f(item, *args, **kwargs)
        Model.save = _s_save
        Model._session_enabled = True

    def _obj_key(self, cls, pk):
        """Returns session key for an object."""
        
        cls_key = getattr(cls, '_cls_key', None)
        if cls_key is None:
            cls_key = cls.__module__ + '.' + cls.__name__
            cls._cls_key = cls_key
        
        return '%s:%s' % (cls_key, pk)

    def add(self, item):
        """
        Adds an object to the session.

        Added object is returned.
        If item's key matches existing key,
        existing object is returned.
        """
        
        if is_persistent(item) is False:
            raise GraphError('Cannot add unpersisted object.')
        
        key = self._obj_key(item.__class__, item.pk)
        existing = self._session_objs.objs.get(key, None)
        if existing is None:
            self._session_objs.objs[key] = item
            return item
        else:
            return existing

    def add_with_dup_check(self, item):
       """Add a persistent object to the session."""

       if self.add(item) is not item:
           raise GraphError('Instance with identical id already exists in session.')