From db1e69813bc0b31957e7589723a2f80cda74454b Mon Sep 17 00:00:00 2001 From: Frank Battaglia Date: Thu, 31 Oct 2013 20:43:07 -0400 Subject: [PATCH] add atomic conditions to save Conflicts: mongoengine/document.py --- mongoengine/document.py | 25 +++++++++++++++++++------ 1 file changed, 19 insertions(+), 6 deletions(-) diff --git a/mongoengine/document.py b/mongoengine/document.py index 114778eb..f870a078 100644 --- a/mongoengine/document.py +++ b/mongoengine/document.py @@ -13,7 +13,8 @@ from mongoengine.base import (DocumentMetaclass, TopLevelDocumentMetaclass, BaseDocument, BaseDict, BaseList, ALLOW_INHERITANCE, get_document) from mongoengine.errors import ValidationError -from mongoengine.queryset import OperationError, NotUniqueError, QuerySet +from mongoengine.queryset import (OperationError, NotUniqueError, + QuerySet, transform) from mongoengine.connection import get_db, DEFAULT_CONNECTION_NAME from mongoengine.context_managers import switch_db, switch_collection @@ -180,7 +181,7 @@ class Document(BaseDocument): def save(self, force_insert=False, validate=True, clean=True, write_concern=None, cascade=None, cascade_kwargs=None, - _refs=None, **kwargs): + _refs=None, save_condition=None, **kwargs): """Save the :class:`~mongoengine.Document` to the database. If the document already exists, it will be updated, otherwise it will be created. @@ -203,7 +204,8 @@ class Document(BaseDocument): :param cascade_kwargs: (optional) kwargs dictionary to be passed throw to cascading saves. Implies ``cascade=True``. :param _refs: A list of processed references used in cascading saves - + :param save_condition: only perform save if matching record in db + satisfies condition(s) (e.g., version number) .. versionchanged:: 0.5 In existing documents it only saves changed fields using set / unset. Saves are cascaded and any @@ -217,6 +219,9 @@ class Document(BaseDocument): meta['cascade'] = True. Also you can pass different kwargs to the cascade save using cascade_kwargs which overwrites the existing kwargs with custom values. + .. versionchanged:: 0.8.5 + Optional save_condition that only overwrites existing documents + if the condition is satisfied in the current db record. """ signals.pre_save.send(self.__class__, document=self) @@ -230,7 +235,8 @@ class Document(BaseDocument): created = ('_id' not in doc or self._created or force_insert) - signals.pre_save_post_validation.send(self.__class__, document=self, created=created) + signals.pre_save_post_validation.send(self.__class__, document=self, + created=created) try: collection = self._get_collection() @@ -243,7 +249,12 @@ class Document(BaseDocument): object_id = doc['_id'] updates, removals = self._delta() # Need to add shard key to query, or you get an error - select_dict = {'_id': object_id} + if save_condition is not None: + select_dict = transform.query(self.__class__, + **save_condition) + else: + select_dict = {} + select_dict['_id'] = object_id shard_key = self.__class__._meta.get('shard_key', tuple()) for k in shard_key: actual_key = self._db_field_map.get(k, k) @@ -263,10 +274,12 @@ class Document(BaseDocument): if removals: update_query["$unset"] = removals if updates or removals: + upsert = save_condition is None last_error = collection.update(select_dict, update_query, - upsert=True, **write_concern) + upsert=upsert, **write_concern) created = is_new_object(last_error) + if cascade is None: cascade = self._meta.get('cascade', False) or cascade_kwargs is not None