From 3cb7e302aaa338681deed81273058c35c4642733 Mon Sep 17 00:00:00 2001
From: Hardy <caoxiaozhen@secnium.com>
Date: Sun, 5 Jul 2020 01:39:42 +0800
Subject: [PATCH] add 'ordered' parameter in QuerySet.bulk_create for duplicate
 insert

---
 pymodm/queryset.py    | 11 ++++++++---
 test/models.py        | 11 +++++++++++
 test/test_queryset.py | 21 +++++++++++++++++++--
 3 files changed, 38 insertions(+), 5 deletions(-)

diff --git a/pymodm/queryset.py b/pymodm/queryset.py
index d7f7be1..2b477c0 100644
--- a/pymodm/queryset.py
+++ b/pymodm/queryset.py
@@ -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 (
@@ -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:
@@ -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`
@@ -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
diff --git a/test/models.py b/test/models.py
index cad4ee2..695939e 100644
--- a/test/models.py
+++ b/test/models.py
@@ -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
 
 
@@ -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)]
diff --git a/test/test_queryset.py b/test/test_queryset.py
index 278cc83..2807d2a 100644
--- a/test/test_queryset.py
+++ b/test/test_queryset.py
@@ -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):
@@ -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)