PYTHON-705 - Fix Bulk API legacy upsert _id compatibility

behackett · behackett · commit 48ffdb0c627f · 2014-06-08T16:38:33.000-07:00
Versions of MongoDB previous to 2.6 only return the upserted
field for an upsert operation if the _id value is an ObjectId.
This patch works around that issue to ensure nUpserted counts
are correct regardless of server version.
diff --git a/pymongo/bulk.py b/pymongo/bulk.py
@@ -103,13 +103,25 @@ def _merge_legacy(run, full_result, result, index):
 
     if run.op_type == _INSERT:
         full_result['nInserted'] += 1
+
     elif run.op_type == _UPDATE:
         if "upserted" in result:
             doc = {"index": run.index(index), "_id": result["upserted"]}
             full_result["upserted"].append(doc)
             full_result['nUpserted'] += affected
+        # Versions of MongoDB before 2.6 don't return the _id for an
+        # upsert if _id is not an ObjectId.
+        elif result.get("updatedExisting") == False and affected == 1:
+            op = run.ops[index]
+            # If _id is in both the update document *and* the query spec
+            # the update document _id takes precedence.
+            _id = op['u'].get('_id', op['q'].get('_id'))
+            doc = {u"index": run.index(index), u"_id": _id}
+            full_result["upserted"].append(doc)
+            full_result['nUpserted'] += affected
         else:
             full_result['nMatched'] += affected
+
     elif run.op_type == _DELETE:
         full_result['nRemoved'] += affected
 
diff --git a/test/test_bulk.py b/test/test_bulk.py
@@ -596,6 +596,28 @@ def test_upsert_large(self):
 
         self.assertEqual(1, self.coll.find({'x': 1}).count())
 
+    def test_client_generated_upsert_id(self):
+        batch = self.coll.initialize_ordered_bulk_op()
+        batch.find({'_id': 0}).upsert().update_one({'$set': {'a': 0}})
+        batch.find({'a': 1}).upsert().replace_one({'_id': 1})
+        if not client_context.version.at_least(2, 6, 0):
+            # This case is only possible in MongoDB versions before 2.6.
+            batch.find({'_id': 3}).upsert().replace_one({'_id': 2})
+        else:
+            # This is just here to make the counts right in all cases.
+            batch.find({'_id': 2}).upsert().replace_one({'_id': 2})
+        result = batch.execute()
+        self.assertEqualResponse(
+            {'nMatched': 0,
+             'nModified': 0,
+             'nUpserted': 3,
+             'nInserted': 0,
+             'nRemoved': 0,
+             'upserted': [{'index': 0, '_id': 0},
+                          {'index': 1, '_id': 1},
+                          {'index': 2, '_id': 2}]},
+            result)
+
     def test_single_ordered_batch(self):
         batch = self.coll.initialize_ordered_bulk_op()
         batch.insert({'a': 1})
