From cd1da3971bd2dff7efd4621aa3e574c98fb8f6d9 Mon Sep 17 00:00:00 2001 From: Jim Vogel Date: Wed, 16 Jun 2021 20:45:15 -0500 Subject: [PATCH 1/2] adding upsert option to save --- mongoengine/document.py | 12 ++++++--- tests/document/test_instance.py | 47 +++++++++++++++++++++++++++++++++ 2 files changed, 55 insertions(+), 4 deletions(-) diff --git a/mongoengine/document.py b/mongoengine/document.py index 0ba5db126..eaf1040c0 100644 --- a/mongoengine/document.py +++ b/mongoengine/document.py @@ -332,6 +332,7 @@ def save( _refs=None, save_condition=None, signal_kwargs=None, + upsert=None, **kwargs, ): """Save the :class:`~mongoengine.Document` to the database. If the @@ -361,6 +362,8 @@ def save( Raises :class:`OperationError` if the conditions are not satisfied :param signal_kwargs: (optional) kwargs dictionary to be passed to the signal calls. + :param upsert: (optional) explicitly forces upsert to value if it is an + update .. versionchanged:: 0.5 In existing documents it only saves changed fields using @@ -407,7 +410,7 @@ def save( object_id = self._save_create(doc, force_insert, write_concern) else: object_id, created = self._save_update( - doc, save_condition, write_concern + doc, save_condition, write_concern, upsert ) if cascade is None: @@ -505,7 +508,7 @@ def _integrate_shard_key(self, doc, select_dict): return select_dict - def _save_update(self, doc, save_condition, write_concern): + def _save_update(self, doc, save_condition, write_concern, upsert=None): """Update an existing document. Helper method, should only be used inside save(). @@ -524,12 +527,13 @@ def _save_update(self, doc, save_condition, write_concern): update_doc = self._get_update_doc() if update_doc: - upsert = save_condition is None + if upsert is None: + upsert = save_condition is None with set_write_concern(collection, write_concern) as wc_collection: last_error = wc_collection.update_one( select_dict, update_doc, upsert=upsert ).raw_result - if not upsert and last_error["n"] == 0: + if save_condition is not None and last_error["n"] == 0: raise SaveConditionError( "Race condition preventing document update detected" ) diff --git a/tests/document/test_instance.py b/tests/document/test_instance.py index 1469c9bb6..13e1c8bab 100644 --- a/tests/document/test_instance.py +++ b/tests/document/test_instance.py @@ -1457,6 +1457,53 @@ def test_inserts_if_you_set_the_pk(self): assert 2 == self.Person.objects.count() + def test_save_upsert_false_doesnt_insert_when_deleted(self): + class Person(Document): + name = StringField() + + Person.drop_collection() + + p1 = Person(name="Wilson Snr") + p1.save() + p2 = Person.objects().first() + p1.delete() + p2.name = " Bob Snr" + p2.save(upsert=False) + + assert Person.objects.count() == 0 + + def test_save_upsert_true_inserts_when_deleted(self): + class Person(Document): + name = StringField() + + Person.drop_collection() + + p1 = Person(name="Wilson Snr") + p1.save() + p2 = Person.objects().first() + p1.delete() + p2.name = "Bob Snr" + p2.save(upsert=True) + + assert Person.objects.count() == 1 + + def test_save_upsert_null_inserts_when_deleted(self): + # probably want to remove this as this is bad but preserved for backwards compatibility + # see https://github.com/MongoEngine/mongoengine/issues/564 + class Person(Document): + name = StringField() + + Person.drop_collection() + + p1 = Person(name="Wilson Snr") + p1.save() + p2 = Person.objects().first() + p1.delete() + p2.name = "Bob Snr" + p2.save(upsert=None) # default if you dont pass it + + assert Person.objects.count() == 1 + def test_can_save_if_not_included(self): class EmbeddedDoc(EmbeddedDocument): pass From 2c807a7a56492e4f2afca44281d6b82c65595417 Mon Sep 17 00:00:00 2001 From: Jim Vogel Date: Sat, 28 Aug 2021 16:11:37 -0500 Subject: [PATCH 2/2] raise exception when upsert is True and a save condition is set --- mongoengine/document.py | 4 ++++ tests/document/test_instance.py | 12 ++++++++++++ 2 files changed, 16 insertions(+) diff --git a/mongoengine/document.py b/mongoengine/document.py index eaf1040c0..184017922 100644 --- a/mongoengine/document.py +++ b/mongoengine/document.py @@ -513,6 +513,10 @@ def _save_update(self, doc, save_condition, write_concern, upsert=None): Helper method, should only be used inside save(). """ + if upsert and (save_condition is not None): + raise ValueError( + "Updating with a save_condition implies upsert is False or None but upsert is True" + ) collection = self._get_collection() object_id = doc["_id"] created = False diff --git a/tests/document/test_instance.py b/tests/document/test_instance.py index 13e1c8bab..60fa62fd9 100644 --- a/tests/document/test_instance.py +++ b/tests/document/test_instance.py @@ -1504,6 +1504,18 @@ class Person(Document): assert Person.objects.count() == 1 + def test_save_upsert_raises_value_error_when_upsert_and_save_condition_set(self): + class Person(Document): + name = StringField() + + Person.drop_collection() + + p1 = Person(name="Wilson Snr") + p1.save() + p1.name = "Bob Snr" + with pytest.raises(ValueError): + p1.save(save_condition={}, upsert=True) + def test_can_save_if_not_included(self): class EmbeddedDoc(EmbeddedDocument): pass