Implementing Write Concern

Added write_options dict to save, update, update_one and get_or_create.
Thanks to justquick for the initial ticket and code.

Refs #132

mongoengine/document.py

@@ -56,7 +56,7 @@ class Document(BaseDocument):
 
     __metaclass__ = TopLevelDocumentMetaclass
 
-    def save(self, safe=True, force_insert=False, validate=True):
+    def save(self, safe=True, force_insert=False, validate=True, write_options=None):
         """Save the :class:`~mongoengine.Document` to the database. If the
         document already exists, it will be updated, otherwise it will be
         created.
@@ -68,16 +68,26 @@ class Document(BaseDocument):
         :param force_insert: only try to create a new document, don't allow
             updates of existing documents
         :param validate: validates the document; set to ``False`` to skip.
+        :param write_options: Extra keyword arguments are passed down to
+                :meth:`~pymongo.collection.Collection.save` OR
+                :meth:`~pymongo.collection.Collection.insert`
+                which will be used as options for the resultant ``getLastError`` command.
+                For example, ``save(..., w=2, fsync=True)`` will wait until at least two servers
+                have recorded the write and will force an fsync on each server being written to.
         """
         if validate:
             self.validate()
+
+        if not write_options:
+            write_options = {}
+
         doc = self.to_mongo()
         try:
             collection = self.__class__.objects._collection
             if force_insert:
-                object_id = collection.insert(doc, safe=safe)
+                object_id = collection.insert(doc, safe=safe, **write_options)
             else:
-                object_id = collection.save(doc, safe=safe)
+                object_id = collection.save(doc, safe=safe, **write_options)
         except pymongo.errors.OperationFailure, err:
             message = 'Could not save document (%s)'
             if u'duplicate key' in unicode(err):

mongoengine/queryset.py

@@ -643,7 +643,7 @@ class QuerySet(object):
             raise self._document.DoesNotExist("%s matching query does not exist."
                                               % self._document._class_name)
 
-    def get_or_create(self, *q_objs, **query):
+    def get_or_create(self, write_options=None, *q_objs, **query):
         """Retrieve unique object or create, if it doesn't exist. Returns a tuple of
         ``(object, created)``, where ``object`` is the retrieved or created object
         and ``created`` is a boolean specifying whether a new object was created. Raises
@@ -653,6 +653,10 @@ class QuerySet(object):
         dictionary of default values for the new document may be provided as a
         keyword argument called :attr:`defaults`.
 
+        :param write_options: optional extra keyword arguments used if we
+            have to create a new document.
+            Passes any write_options onto :meth:`~mongoengine.document.Document.save`
+
         .. versionadded:: 0.3
         """
         defaults = query.get('defaults', {})
@@ -664,7 +668,7 @@ class QuerySet(object):
         if count == 0:
             query.update(defaults)
             doc = self._document(**query)
-            doc.save()
+            doc.save(write_options=write_options)
             return doc, True
         elif count == 1:
             return self.first(), False
@@ -1055,22 +1059,27 @@ class QuerySet(object):
 
         return mongo_update
 
-    def update(self, safe_update=True, upsert=False, **update):
+    def update(self, safe_update=True, upsert=False, write_options=None, **update):
         """Perform an atomic update on the fields matched by the query. When
         ``safe_update`` is used, the number of affected documents is returned.
 
-        :param safe: check if the operation succeeded before returning
+        :param safe_update: check if the operation succeeded before returning
-        :param update: Django-style update keyword arguments
+        :param upsert: Any existing document with that "_id" is overwritten.
+        :param write_options: extra keyword arguments for :meth:`~pymongo.collection.Collection.update`
 
         .. versionadded:: 0.2
         """
         if pymongo.version < '1.1.1':
             raise OperationError('update() method requires PyMongo 1.1.1+')
 
+        if not write_options:
+            write_options = {}
+
         update = QuerySet._transform_update(self._document, **update)
         try:
             ret = self._collection.update(self._query, update, multi=True,
-                                          upsert=upsert, safe=safe_update)
+                                          upsert=upsert, safe=safe_update,
+                                          **write_options)
             if ret is not None and 'n' in ret:
                 return ret['n']
         except pymongo.errors.OperationFailure, err:
@@ -1079,22 +1088,27 @@ class QuerySet(object):
                 raise OperationError(message)
             raise OperationError(u'Update failed (%s)' % unicode(err))
 
-    def update_one(self, safe_update=True, upsert=False, **update):
+    def update_one(self, safe_update=True, upsert=False, write_options=None, **update):
         """Perform an atomic update on first field matched by the query. When
         ``safe_update`` is used, the number of affected documents is returned.
 
-        :param safe: check if the operation succeeded before returning
+        :param safe_update: check if the operation succeeded before returning
+        :param upsert: Any existing document with that "_id" is overwritten.
+        :param write_options: extra keyword arguments for :meth:`~pymongo.collection.Collection.update`
         :param update: Django-style update keyword arguments
 
         .. versionadded:: 0.2
         """
+        if not write_options:
+            write_options = {}
         update = QuerySet._transform_update(self._document, **update)
         try:
             # Explicitly provide 'multi=False' to newer versions of PyMongo
             # as the default may change to 'True'
             if pymongo.version >= '1.1.1':
                 ret = self._collection.update(self._query, update, multi=False,
-                                              upsert=upsert, safe=safe_update)
+                                              upsert=upsert, safe=safe_update,
+                                              **write_options)
             else:
                 # Older versions of PyMongo don't support 'multi'
                 ret = self._collection.update(self._query, update,

tests/document.py

@@ -898,6 +898,26 @@ class DocumentTest(unittest.TestCase):
         pickle_doc.reload()
         self.assertEquals(resurrected, pickle_doc)
 
+    def test_write_options(self):
+        """Test that passing write_options works"""
+
+        self.Person.drop_collection()
+
+        write_options = {"fsync": True}
+
+        author, created = self.Person.objects.get_or_create(
+                            name='Test User', write_options=write_options)
+        author.save(write_options=write_options)
+
+        self.Person.objects.update(set__name='Ross', write_options=write_options)
+
+        author = self.Person.objects.first()
+        self.assertEquals(author.name, 'Ross')
+
+        self.Person.objects.update_one(set__name='Test User', write_options=write_options)
+        author = self.Person.objects.first()
+        self.assertEquals(author.name, 'Test User')
+
 
 if __name__ == '__main__':
     unittest.main()