PYTHON-1336 Add collation support to Collection.bulk_write.

ShaneHarvey · ShaneHarvey · commit 8eb1aefeff9b · 2017-08-03T15:21:04.000-07:00
diff --git a/doc/changelog.rst b/doc/changelog.rst
@@ -17,6 +17,11 @@ Highlights:
   overhead of :class:`~pymongo.operations.InsertOne`,
   :class:`~pymongo.operations.DeleteOne`, and
   :class:`~pymongo.operations.DeleteMany`.
+- Added the `collation` option to :class:`~pymongo.operations.DeleteOne`,
+  :class:`~pymongo.operations.DeleteMany`,
+  :class:`~pymongo.operations.ReplaceOne`,
+  :class:`~pymongo.operations.UpdateOne`, and
+  :class:`~pymongo.operations.UpdateMany`.
 
 Changes and Deprecations:
 
diff --git a/pymongo/bulk.py b/pymongo/bulk.py
@@ -635,8 +635,8 @@ def find(self, selector, collation=None):
           - `selector` (dict): the selection criteria for update
             and remove operations.
           - `collation` (optional): An instance of
-            :class:`~pymongo.collation.Collation`. This option is only supported
-            on MongoDB 3.4 and above.
+            :class:`~pymongo.collation.Collation`. This option is only
+            supported on MongoDB 3.4 and above.
 
         :Returns:
           - A :class:`BulkWriteOperation` instance, used to add
diff --git a/pymongo/operations.py b/pymongo/operations.py
@@ -54,30 +54,38 @@ def __ne__(self, other):
 class DeleteOne(object):
     """Represents a delete_one operation."""
 
-    __slots__ = ("_filter",)
+    __slots__ = ("_filter", "_collation")
 
-    def __init__(self, filter):
+    def __init__(self, filter, collation=None):
         """Create a DeleteOne instance.
 
         For use with :meth:`~pymongo.collection.Collection.bulk_write`.
 
         :Parameters:
           - `filter`: A query that matches the document to delete.
+          - `collation` (optional): An instance of
+            :class:`~pymongo.collation.Collation`. This option is only
+            supported on MongoDB 3.4 and above.
+
+        .. versionchanged:: 3.5
+           Added the `collation` option.
         """
         if filter is not None:
             validate_is_mapping("filter", filter)
         self._filter = filter
+        self._collation = collation
 
     def _add_to_bulk(self, bulkobj):
         """Add this operation to the _Bulk instance `bulkobj`."""
-        bulkobj.add_delete(self._filter, 1)
+        bulkobj.add_delete(self._filter, 1, collation=self._collation)
 
     def __repr__(self):
-        return "DeleteOne(%r)" % (self._filter,)
+        return "DeleteOne(%r, %r)" % (self._filter, self._collation)
 
     def __eq__(self, other):
         if type(other) == type(self):
-            return other._filter == self._filter
+            return ((other._filter, other._collation) ==
+                    (self._filter, self._collation))
         return NotImplemented
 
     def __ne__(self, other):
@@ -87,30 +95,38 @@ def __ne__(self, other):
 class DeleteMany(object):
     """Represents a delete_many operation."""
 
-    __slots__ = ("_filter",)
+    __slots__ = ("_filter", "_collation")
 
-    def __init__(self, filter):
+    def __init__(self, filter, collation=None):
         """Create a DeleteMany instance.
 
         For use with :meth:`~pymongo.collection.Collection.bulk_write`.
 
         :Parameters:
           - `filter`: A query that matches the documents to delete.
+          - `collation` (optional): An instance of
+            :class:`~pymongo.collation.Collation`. This option is only
+            supported on MongoDB 3.4 and above.
+
+        .. versionchanged:: 3.5
+           Added the `collation` option.
         """
         if filter is not None:
             validate_is_mapping("filter", filter)
         self._filter = filter
+        self._collation = collation
 
     def _add_to_bulk(self, bulkobj):
         """Add this operation to the _Bulk instance `bulkobj`."""
-        bulkobj.add_delete(self._filter, 0)
+        bulkobj.add_delete(self._filter, 0, collation=self._collation)
 
     def __repr__(self):
-        return "DeleteMany(%r)" % (self._filter,)
+        return "DeleteMany(%r, %r)" % (self._filter, self._collation)
 
     def __eq__(self, other):
         if type(other) == type(self):
-            return other._filter == self._filter
+            return ((other._filter, other._collation) ==
+                    (self._filter, self._collation))
         return NotImplemented
 
     def __ne__(self, other):
@@ -120,37 +136,40 @@ def __ne__(self, other):
 class _UpdateOp(object):
     """Private base class for update operations."""
 
-    __slots__ = ("_filter", "_doc", "_upsert")
+    __slots__ = ("_filter", "_doc", "_upsert", "_collation")
 
-    def __init__(self, filter, doc, upsert):
+    def __init__(self, filter, doc, upsert, collation):
         if filter is not None:
             validate_is_mapping("filter", filter)
         if upsert is not None:
             validate_boolean("upsert", upsert)
         self._filter = filter
         self._doc = doc
         self._upsert = upsert
+        self._collation = collation
 
     def __eq__(self, other):
         if type(other) == type(self):
-            return ((other._filter, other._doc, other._upsert) ==
-                    (self._filter, self._doc, self._upsert))
+            return (
+                (other._filter, other._doc, other._upsert, other._collation) ==
+                (self._filter, self._doc, self._upsert, self._collation))
         return NotImplemented
 
     def __ne__(self, other):
         return not self == other
 
     def __repr__(self):
-        return "%s(%r, %r, %r)" % (
-            self.__class__.__name__, self._filter, self._doc, self._upsert)
+        return "%s(%r, %r, %r, %r)" % (
+            self.__class__.__name__, self._filter, self._doc, self._upsert,
+            self._collation)
 
 
 class ReplaceOne(_UpdateOp):
     """Represents a replace_one operation."""
 
     __slots__ = ()
 
-    def __init__(self, filter, replacement, upsert=False):
+    def __init__(self, filter, replacement, upsert=False, collation=None):
         """Create a ReplaceOne instance.
 
         For use with :meth:`~pymongo.collection.Collection.bulk_write`.
@@ -160,20 +179,28 @@ def __init__(self, filter, replacement, upsert=False):
           - `replacement`: The new document.
           - `upsert` (optional): If ``True``, perform an insert if no documents
             match the filter.
+          - `collation` (optional): An instance of
+            :class:`~pymongo.collation.Collation`. This option is only
+            supported on MongoDB 3.4 and above.
+
+        .. versionchanged:: 3.5
+           Added the `collation` option.
         """
-        super(ReplaceOne, self).__init__(filter, replacement, upsert)
+        super(ReplaceOne, self).__init__(filter, replacement, upsert,
+                                         collation)
 
     def _add_to_bulk(self, bulkobj):
         """Add this operation to the _Bulk instance `bulkobj`."""
-        bulkobj.add_replace(self._filter, self._doc, self._upsert)
+        bulkobj.add_replace(self._filter, self._doc, self._upsert,
+                            collation=self._collation)
 
 
 class UpdateOne(_UpdateOp):
     """Represents an update_one operation."""
 
     __slots__ = ()
 
-    def __init__(self, filter, update, upsert=False):
+    def __init__(self, filter, update, upsert=False, collation=None):
         """Represents an update_one operation.
 
         For use with :meth:`~pymongo.collection.Collection.bulk_write`.
@@ -183,20 +210,27 @@ def __init__(self, filter, update, upsert=False):
           - `update`: The modifications to apply.
           - `upsert` (optional): If ``True``, perform an insert if no documents
             match the filter.
+          - `collation` (optional): An instance of
+            :class:`~pymongo.collation.Collation`. This option is only
+            supported on MongoDB 3.4 and above.
+
+        .. versionchanged:: 3.5
+           Added the `collation` option.
         """
-        super(UpdateOne, self).__init__(filter, update, upsert)
+        super(UpdateOne, self).__init__(filter, update, upsert, collation)
 
     def _add_to_bulk(self, bulkobj):
         """Add this operation to the _Bulk instance `bulkobj`."""
-        bulkobj.add_update(self._filter, self._doc, False, self._upsert)
+        bulkobj.add_update(self._filter, self._doc, False, self._upsert,
+                           collation=self._collation)
 
 
 class UpdateMany(_UpdateOp):
     """Represents an update_many operation."""
 
     __slots__ = ()
 
-    def __init__(self, filter, update, upsert=False):
+    def __init__(self, filter, update, upsert=False, collation=None):
         """Create an UpdateMany instance.
 
         For use with :meth:`~pymongo.collection.Collection.bulk_write`.
@@ -206,12 +240,19 @@ def __init__(self, filter, update, upsert=False):
           - `update`: The modifications to apply.
           - `upsert` (optional): If ``True``, perform an insert if no documents
             match the filter.
+          - `collation` (optional): An instance of
+            :class:`~pymongo.collation.Collation`. This option is only
+            supported on MongoDB 3.4 and above.
+
+        .. versionchanged:: 3.5
+           Added the `collation` option.
         """
-        super(UpdateMany, self).__init__(filter, update, upsert)
+        super(UpdateMany, self).__init__(filter, update, upsert, collation)
 
     def _add_to_bulk(self, bulkobj):
         """Add this operation to the _Bulk instance `bulkobj`."""
-        bulkobj.add_update(self._filter, self._doc, True, self._upsert)
+        bulkobj.add_update(self._filter, self._doc, True, self._upsert,
+                           collation=self._collation)
 
 
 class IndexModel(object):
diff --git a/test/test_collation.py b/test/test_collation.py
@@ -27,7 +27,8 @@
     CollationMaxVariable)
 from pymongo.errors import ConfigurationError
 from pymongo.write_concern import WriteConcern
-from pymongo.operations import IndexModel
+from pymongo.operations import (DeleteMany, DeleteOne, IndexModel, ReplaceOne,
+                                UpdateMany, UpdateOne)
 
 
 class TestCollationObject(unittest.TestCase):
@@ -273,35 +274,75 @@ def test_find_and(self):
                                           collation=self.collation)
         self.assertCollationInLastCommand()
 
+    @raisesConfigurationErrorForOldMongoDB
+    def test_bulk_write(self):
+        self.db.test.collection.bulk_write([
+            DeleteOne({'noCollation': 42}),
+            DeleteMany({'noCollation': 42}),
+            DeleteOne({'foo': 42}, collation=self.collation),
+            DeleteMany({'foo': 42}, collation=self.collation),
+            ReplaceOne({'noCollation': 24}, {'bar': 42}),
+            UpdateOne({'noCollation': 84}, {'$set': {'bar': 10}}, upsert=True),
+            UpdateMany({'noCollation': 45}, {'$set': {'bar': 42}}),
+            ReplaceOne({'foo': 24}, {'foo': 42}, collation=self.collation),
+            UpdateOne({'foo': 84}, {'$set': {'foo': 10}}, upsert=True,
+                      collation=self.collation),
+            UpdateMany({'foo': 45}, {'$set': {'foo': 42}},
+                       collation=self.collation)
+        ])
+
+        delete_cmd = self.listener.results['started'][0].command
+        update_cmd = self.listener.results['started'][1].command
+
+        def check_ops(ops):
+            for op in ops:
+                if 'noCollation' in op['q']:
+                    self.assertNotIn('collation', op)
+                else:
+                    self.assertEqual(self.collation.document,
+                                     op['collation'])
+
+        check_ops(delete_cmd['deletes'])
+        check_ops(update_cmd['updates'])
+
     @raisesConfigurationErrorForOldMongoDB
     def test_bulk(self):
         bulk = self.db.test.initialize_ordered_bulk_op()
+        bulk.find({'noCollation': 42}).remove_one()
+        bulk.find({'noCollation': 42}).remove()
         bulk.find({'foo': 42}, collation=self.collation).remove_one()
-        # No collation for this operation.
-        bulk.find({'bar': 24}).replace_one({'bar': 42})
-        bulk.find({'foo': 84}, collation=self.collation).upsert()
+        bulk.find({'foo': 42}, collation=self.collation).remove()
+        bulk.find({'noCollation': 24}).replace_one({'bar': 42})
+        bulk.find({'noCollation': 84}).upsert().update_one(
+            {'$set': {'foo': 10}})
+        bulk.find({'noCollation': 45}).update({'$set': {'bar': 42}})
+        bulk.find({'foo': 24}, collation=self.collation).replace_one(
+            {'foo': 42})
+        bulk.find({'foo': 84}, collation=self.collation).upsert().update_one(
+            {'$set': {'foo': 10}})
         bulk.find({'foo': 45}, collation=self.collation).update({
             '$set': {'foo': 42}})
         bulk.execute()
 
         delete_cmd = self.listener.results['started'][0].command
         update_cmd = self.listener.results['started'][1].command
 
-        for op in delete_cmd['deletes']:
-            self.assertEqual(self.collation.document,
-                             op['collation'])
+        def check_ops(ops):
+            for op in ops:
+                if 'noCollation' in op['q']:
+                    self.assertNotIn('collation', op)
+                else:
+                    self.assertEqual(self.collation.document,
+                                     op['collation'])
 
-        for op in update_cmd['updates']:
-            if 'foo' in op['q']:
-                self.assertEqual(self.collation.document,
-                                 op['collation'])
-            else:
-                self.assertNotIn('collation', op)
+        check_ops(delete_cmd['deletes'])
+        check_ops(update_cmd['updates'])
 
     @client_context.require_version_max(3, 3, 8)
     def test_mixed_bulk_collation(self):
         bulk = self.db.test.initialize_unordered_bulk_op()
-        bulk.find({'foo': 42}).upsert()
+        bulk.find({'foo': 42}).upsert().update_one(
+            {'$set': {'bar': 10}})
         bulk.find({'foo': 43}, collation=self.collation).remove_one()
         with self.assertRaises(ConfigurationError):
             bulk.execute()
@@ -343,6 +384,10 @@ def test_unacknowledged_write(self):
             {'$set': {'hello': 'moon'}})
         with self.assertRaises(ConfigurationError):
             bulk.execute()
+        update_one = UpdateOne({'hello': 'world'}, {'$set': {'hello': 'moon'}},
+                               collation=self.collation)
+        with self.assertRaises(ConfigurationError):
+            collection.bulk_write([update_one])
 
     @raisesConfigurationErrorForOldMongoDB
     def test_cursor_collation(self):
