From 0bb9781b91bd3e70c84d1bdfe8bb1b85baaab70b Mon Sep 17 00:00:00 2001 From: Greg Banks Date: Thu, 26 Apr 2012 13:56:52 -0700 Subject: [PATCH 1/2] add "safe" and "write_options" parameters to QuerySet.insert similar to Document.save --- mongoengine/document.py | 5 +++-- mongoengine/queryset.py | 24 ++++++++++++++++++++++-- tests/queryset.py | 21 +++++++++++++++++++-- 3 files changed, 44 insertions(+), 6 deletions(-) diff --git a/mongoengine/document.py b/mongoengine/document.py index 4b5506f4..f4c74c7e 100644 --- a/mongoengine/document.py +++ b/mongoengine/document.py @@ -147,8 +147,9 @@ class Document(BaseDocument): :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. + For example, ``save(..., write_options={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. :param cascade: Sets the flag for cascading saves. You can set a default by setting "cascade" in the document __meta__ :param cascade_kwargs: optional kwargs dictionary to be passed throw to cascading saves diff --git a/mongoengine/queryset.py b/mongoengine/queryset.py index 6d1c9c13..98758226 100644 --- a/mongoengine/queryset.py +++ b/mongoengine/queryset.py @@ -824,11 +824,21 @@ class QuerySet(object): result = None return result - def insert(self, doc_or_docs, load_bulk=True): + def insert(self, doc_or_docs, load_bulk=True, safe=False, write_options=None): """bulk insert documents + If ``safe=True`` and the operation is unsuccessful, an + :class:`~mongoengine.OperationError` will be raised. + :param docs_or_doc: a document or list of documents to be inserted :param load_bulk (optional): If True returns the list of document instances + :param safe: check if the operation succeeded before returning + :param write_options: Extra keyword arguments are passed down to + :meth:`~pymongo.collection.Collection.insert` + which will be used as options for the resultant ``getLastError`` command. + For example, ``insert(..., {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. By default returns document instances, set ``load_bulk`` to False to return just ``ObjectIds`` @@ -837,6 +847,10 @@ class QuerySet(object): """ from document import Document + if not write_options: + write_options = {} + write_options.update({'safe': safe}) + docs = doc_or_docs return_one = False if isinstance(docs, Document) or issubclass(docs.__class__, Document): @@ -854,7 +868,13 @@ class QuerySet(object): raw.append(doc.to_mongo()) signals.pre_bulk_insert.send(self._document, documents=docs) - ids = self._collection.insert(raw) + try: + ids = self._collection.insert(raw, **write_options) + except pymongo.errors.OperationFailure, err: + message = 'Could not save document (%s)' + if u'duplicate key' in unicode(err): + message = u'Tried to save duplicate unique keys (%s)' + raise OperationError(message % unicode(err)) if not load_bulk: signals.post_bulk_insert.send( diff --git a/tests/queryset.py b/tests/queryset.py index 0d5aaabb..4cf11657 100644 --- a/tests/queryset.py +++ b/tests/queryset.py @@ -480,7 +480,7 @@ class QuerySetTest(unittest.TestCase): self.assertEqual(person.name, "User C") def test_bulk_insert(self): - """Ensure that query by array position works. + """Ensure that bulk insert works """ class Comment(EmbeddedDocument): @@ -490,7 +490,7 @@ class QuerySetTest(unittest.TestCase): comments = ListField(EmbeddedDocumentField(Comment)) class Blog(Document): - title = StringField() + title = StringField(unique=True) tags = ListField(StringField()) posts = ListField(EmbeddedDocumentField(Post)) @@ -563,6 +563,23 @@ class QuerySetTest(unittest.TestCase): obj_id = Blog.objects.insert(blog1, load_bulk=False) self.assertEquals(obj_id.__class__.__name__, 'ObjectId') + Blog.drop_collection() + post3 = Post(comments=[comment1, comment1]) + blog1 = Blog(title="foo", posts=[post1, post2]) + blog2 = Blog(title="bar", posts=[post2, post3]) + blog3 = Blog(title="baz", posts=[post1, post2]) + Blog.objects.insert([blog1, blog2]) + + def throw_operation_error_not_unique(): + Blog.objects.insert([blog2, blog3], safe=True) + + self.assertRaises(OperationError, throw_operation_error_not_unique) + self.assertEqual(Blog.objects.count(), 2) + + Blog.objects.insert([blog2, blog3], write_options={'continue_on_error': True}) + self.assertEqual(Blog.objects.count(), 3) + + def test_slave_okay(self): """Ensures that a query can take slave_okay syntax """ From 410443471cc545c48e890e083b49cddbf464030a Mon Sep 17 00:00:00 2001 From: Greg Banks Date: Thu, 26 Apr 2012 14:03:30 -0700 Subject: [PATCH 2/2] TopLevelDocumentMetaClass sets _meta['index_opts'] not _meta['index_options'] --- mongoengine/queryset.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mongoengine/queryset.py b/mongoengine/queryset.py index 98758226..23d7515b 100644 --- a/mongoengine/queryset.py +++ b/mongoengine/queryset.py @@ -498,7 +498,7 @@ class QuerySet(object): background = self._document._meta.get('index_background', False) drop_dups = self._document._meta.get('index_drop_dups', False) - index_opts = self._document._meta.get('index_options', {}) + index_opts = self._document._meta.get('index_opts', {}) index_types = self._document._meta.get('index_types', True) # determine if an index which we are creating includes