Added switch_collection context manager and method (#220)

mongoengine/context_managers.py

@@ -2,7 +2,7 @@ from mongoengine.common import _import_class
 from mongoengine.connection import DEFAULT_CONNECTION_NAME, get_db
 from mongoengine.queryset import OperationError, QuerySet
 
-__all__ = ("switch_db", "no_dereference", "query_counter")
+__all__ = ("switch_db", "switch_collection", "no_dereference", "query_counter")
 
 
 class switch_db(object):
@@ -47,6 +47,49 @@ class switch_db(object):
         self.cls._collection = self.collection
 
 
+class switch_collection(object):
+    """ switch_collection alias context manager.
+
+    Example ::
+
+        class Group(Document):
+            name = StringField()
+
+        Group(name="test").save()  # Saves in the default db
+
+        with switch_collection(Group, 'group1') as Group:
+            Group(name="hello testdb!").save()  # Saves in group1 collection
+
+    """
+
+    def __init__(self, cls, collection_name):
+        """ Construct the switch_collection context manager
+
+        :param cls: the class to change the registered db
+        :param collection_name: the name of the collection to use
+        """
+        self.cls = cls
+        self.ori_collection = cls._get_collection()
+        self.ori_get_collection_name = cls._get_collection_name
+        self.collection_name = collection_name
+
+    def __enter__(self):
+        """ change the _get_collection_name and clear the cached collection """
+
+        @classmethod
+        def _get_collection_name(cls):
+            return self.collection_name
+
+        self.cls._get_collection_name = _get_collection_name
+        self.cls._collection = None
+        return self.cls
+
+    def __exit__(self, t, value, traceback):
+        """ Reset the collection """
+        self.cls._collection = self.ori_collection
+        self.cls._get_collection_name = self.ori_get_collection_name
+
+
 class no_dereference(object):
     """ no_dereference context manager.
 

mongoengine/document.py

@@ -11,7 +11,7 @@ from mongoengine.base import (DocumentMetaclass, TopLevelDocumentMetaclass,
                               ALLOW_INHERITANCE, get_document)
 from mongoengine.queryset import OperationError, NotUniqueError
 from mongoengine.connection import get_db, DEFAULT_CONNECTION_NAME
-from mongoengine.context_managers import switch_db
+from mongoengine.context_managers import switch_db, switch_collection
 
 __all__ = ('Document', 'EmbeddedDocument', 'DynamicDocument',
            'DynamicEmbeddedDocument', 'OperationError',
@@ -398,6 +398,31 @@ class Document(BaseDocument):
         self._objects._collection_obj = collection
         return self
 
+    def switch_collection(self, collection_name):
+        """
+        Temporarily switch the collection for a document instance.
+
+        Only really useful for archiving off data and calling `save()`::
+
+            user = User.objects.get(id=user_id)
+            user.switch_collection('old-users')
+            user.save()
+
+        If you need to read from another database see
+        :class:`~mongoengine.context_managers.switch_collection`
+
+        :param collection_name: The database alias to use for saving the
+            document
+        """
+        with switch_collection(self.__class__, collection_name) as cls:
+            collection = cls._get_collection()
+        self._get_collection = lambda: collection
+        self._collection = collection
+        self._created = True
+        self._objects = self.__class__.objects
+        self._objects._collection_obj = collection
+        return self
+
     def select_related(self, max_depth=1):
         """Handles dereferencing of :class:`~bson.dbref.DBRef` objects to
         a maximum depth in order to cut down the number queries to mongodb.