From 40b69baa2991f750bdf879058ddb2d1e6a9a0c05 Mon Sep 17 00:00:00 2001 From: Ross Lawley Date: Thu, 19 May 2011 16:49:00 +0100 Subject: [PATCH] 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 | 38 ++++++++++++++++++++++++-------------- mongoengine/queryset.py | 32 +++++++++++++++++++++++--------- tests/document.py | 20 ++++++++++++++++++++ 3 files changed, 67 insertions(+), 23 deletions(-) diff --git a/mongoengine/document.py b/mongoengine/document.py index 196662c3..771b9229 100644 --- a/mongoengine/document.py +++ b/mongoengine/document.py @@ -40,44 +40,54 @@ class Document(BaseDocument): presence of `_cls` and `_types`, set :attr:`allow_inheritance` to ``False`` in the :attr:`meta` dictionary. - A :class:`~mongoengine.Document` may use a **Capped Collection** by + A :class:`~mongoengine.Document` may use a **Capped Collection** by specifying :attr:`max_documents` and :attr:`max_size` in the :attr:`meta` dictionary. :attr:`max_documents` is the maximum number of documents that - is allowed to be stored in the collection, and :attr:`max_size` is the - maximum size of the collection in bytes. If :attr:`max_size` is not - specified and :attr:`max_documents` is, :attr:`max_size` defaults to + is allowed to be stored in the collection, and :attr:`max_size` is the + maximum size of the collection in bytes. If :attr:`max_size` is not + specified and :attr:`max_documents` is, :attr:`max_size` defaults to 10000000 bytes (10MB). Indexes may be created by specifying :attr:`indexes` in the :attr:`meta` - dictionary. The value should be a list of field names or tuples of field + dictionary. The value should be a list of field names or tuples of field names. Index direction may be specified by prefixing the field names with a **+** or **-** sign. """ __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. - If ``safe=True`` and the operation is unsuccessful, an + If ``safe=True`` and the operation is unsuccessful, an :class:`~mongoengine.OperationError` will be raised. :param safe: check if the operation succeeded before returning - :param force_insert: only try to create a new document, don't allow + :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): @@ -131,9 +141,9 @@ class MapReduceDocument(object): """A document returned from a map/reduce query. :param collection: An instance of :class:`~pymongo.Collection` - :param key: Document/result key, often an instance of - :class:`~pymongo.objectid.ObjectId`. If supplied as - an ``ObjectId`` found in the given ``collection``, + :param key: Document/result key, often an instance of + :class:`~pymongo.objectid.ObjectId`. If supplied as + an ``ObjectId`` found in the given ``collection``, the object can be accessed via the ``object`` property. :param value: The result(s) for this key. @@ -148,7 +158,7 @@ class MapReduceDocument(object): @property def object(self): - """Lazy-load the object referenced by ``self.key``. ``self.key`` + """Lazy-load the object referenced by ``self.key``. ``self.key`` should be the ``primary_key``. """ id_field = self._document()._meta['id_field'] diff --git a/mongoengine/queryset.py b/mongoengine/queryset.py index 6da11fa7..683aac50 100644 --- a/mongoengine/queryset.py +++ b/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 update: Django-style update keyword arguments + :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` .. 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, diff --git a/tests/document.py b/tests/document.py index 84d0068b..cef6e8c1 100644 --- a/tests/document.py +++ b/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()