Skip to content
This repository was archived by the owner on Jan 21, 2025. It is now read-only.

add 'ordered' parameter in QuerySet.bulk_create for duplicate insert #68

Open
wants to merge 1 commit into
base: master
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 8 additions & 3 deletions pymodm/queryset.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,8 +14,8 @@

import copy

from bson.son import SON
import pymongo
from bson.son import SON

from pymodm import errors
from pymodm.common import (
Expand Down Expand Up @@ -390,7 +390,7 @@ def create(self, **kwargs):
"""
return self._model(**kwargs).save()

def bulk_create(self, object_or_objects, retrieve=False, full_clean=False):
def bulk_create(self, object_or_objects, retrieve=False, full_clean=False, ordered=True):
"""Save Model instances in bulk.

:parameters:
Expand All @@ -402,6 +402,11 @@ def bulk_create(self, object_or_objects, retrieve=False, full_clean=False):
- `full_clean`: Whether to validate each object by calling
the :meth:`~pymodm.MongoModel.full_clean` method before saving.
This isn't done by default.
- `ordered` (optional): If ``True`` (the default) documents will be
inserted on the server serially, in the order provided. If an error
occurs all remaining inserts are aborted. If ``False``, documents
will be inserted on the server in arbitrary order, possibly in
parallel, and all document inserts will be attempted.

:returns: A list of ids for the documents saved, or of the
:class:`~pymodm.MongoModel` instances themselves if `retrieve`
Expand All @@ -428,7 +433,7 @@ def bulk_create(self, object_or_objects, retrieve=False, full_clean=False):
for object in object_or_objects:
object.full_clean()
docs = (obj.to_son() for obj in object_or_objects)
ids = self._collection.insert_many(docs).inserted_ids
ids = self._collection.insert_many(docs, ordered=ordered).inserted_ids
if retrieve:
return list(self.raw({'_id': {'$in': ids}}))
return ids
Expand Down
11 changes: 11 additions & 0 deletions test/models.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,8 @@
# See the License for the specific language governing permissions and
# limitations under the License.

from pymongo import IndexModel

from pymodm import MongoModel, fields


Expand All @@ -33,3 +35,12 @@ class Meta:

class User(ParentModel):
address = fields.CharField()


class Fruit(MongoModel):
name = fields.CharField()
color = fields.CharField()

class Meta:
collection_name = 'test_fruit'
indexes = [IndexModel([('name', 1), ('color', 1)], unique=True)]
21 changes: 19 additions & 2 deletions test/test_queryset.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,14 +13,14 @@
# limitations under the License.

from bson.objectid import ObjectId
from pymongo.errors import BulkWriteError

from pymodm import fields
from pymodm.base import MongoModel
from pymodm.compat import text_type
from pymodm.context_managers import no_auto_dereference

from test import ODMTestCase
from test.models import ParentModel, User
from test.models import ParentModel, User, Fruit


class Vacation(MongoModel):
Expand Down Expand Up @@ -246,3 +246,20 @@ class Post(MongoModel):
posts = list(Post.objects.select_related())
self.assertEqual([], posts[0].comments)
self.assertEqual(posts[1].comments, comments)


class QuerySetBulkCreateTestCase(ODMTestCase):
def setUp(self):
pass

def test_bulk_create_duplicate_key(self):
with self.assertRaises(BulkWriteError) as context:
Fruit.objects.bulk_create([
Fruit(name="Apple", color="Green"),
Fruit(name="Apple", color="Red"),
Fruit(name="Apple", color="Red"),
Fruit(name="Orange", color="Green"),
Fruit(name="Orange", color="Orange"),
Fruit(name="Orange", color="Orange"),
], ordered=False)
self.assertTrue('batch op errors occurred' in context.exception)