From f27debe7f974449f34b0acc9a0c8a8de7f1be437 Mon Sep 17 00:00:00 2001 From: Dmitry Balabanov Date: Thu, 30 Aug 2012 12:40:44 +0400 Subject: [PATCH] Respect sharding key when delete object from collection --- mongoengine/document.py | 18 ++++++++++++------ 1 file changed, 12 insertions(+), 6 deletions(-) diff --git a/mongoengine/document.py b/mongoengine/document.py index f8bf769d..1af7d017 100644 --- a/mongoengine/document.py +++ b/mongoengine/document.py @@ -265,6 +265,16 @@ class Document(BaseDocument): ref.save(**kwargs) ref._changed_fields = [] + @property + def _object_key(self): + """Dict to identify object in collection + """ + select_dict = {'pk': self.pk} + shard_key = self.__class__._meta.get('shard_key', tuple()) + for k in shard_key: + select_dict[k] = getattr(self, k) + return select_dict + def update(self, **kwargs): """Performs an update on the :class:`~mongoengine.Document` A convenience wrapper to :meth:`~mongoengine.QuerySet.update`. @@ -276,11 +286,7 @@ class Document(BaseDocument): raise OperationError('attempt to update a document not yet saved') # Need to add shard key to query, or you get an error - select_dict = {'pk': self.pk} - shard_key = self.__class__._meta.get('shard_key', tuple()) - for k in shard_key: - select_dict[k] = getattr(self, k) - return self.__class__.objects(**select_dict).update_one(**kwargs) + return self.__class__.objects(**self._object_key).update_one(**kwargs) def delete(self, safe=False): """Delete the :class:`~mongoengine.Document` from the database. This @@ -291,7 +297,7 @@ class Document(BaseDocument): signals.pre_delete.send(self.__class__, document=self) try: - self.__class__.objects(pk=self.pk).delete(safe=safe) + self.__class__.objects(**self._object_key).delete(safe=safe) except pymongo.errors.OperationFailure, err: message = u'Could not delete document (%s)' % err.message raise OperationError(message)