Added switch_collection context manager and method (#220)

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.