adding DuplicateKeyError for safe inserts/updates

Author: Mike Dirolf

pymongo/collection.py

@@ -23,7 +23,7 @@
 from objectid import ObjectId
 from cursor import Cursor
 from son import SON
-from errors import InvalidName, OperationFailure
+from errors import InvalidName
 from code import Code
 
 _ZERO = "\x00\x00\x00\x00"

pymongo/connection.py

@@ -44,7 +44,7 @@
 import warnings
 
 from errors import ConnectionFailure, ConfigurationError, AutoReconnect
-from errors import OperationFailure
+from errors import OperationFailure, DuplicateKeyError
 from database import Database
 from cursor_manager import CursorManager
 import bson
@@ -441,7 +441,11 @@ def __check_response_to_last_error(self, response):
             return
         if error["err"] == "not master":
             self._reset()
-        raise OperationFailure(error["err"])
+
+        if "code" in error and error["code"] in [11000, 11001]:
+            raise DuplicateKeyError(error["err"])
+        else:
+            raise OperationFailure(error["err"])
 
     def _send_message(self, message, with_last_error=False):
         """Say something to Mongo.

pymongo/cursor.py

@@ -22,7 +22,7 @@
 import message
 from son import SON
 from code import Code
-from errors import InvalidOperation, OperationFailure, AutoReconnect
+from errors import InvalidOperation, AutoReconnect
 
 _QUERY_OPTIONS = {
     "tailable_cursor": 2,

pymongo/errors.py

@@ -42,6 +42,11 @@ class OperationFailure(Exception):
     """
 
 
+class DuplicateKeyError(OperationFailure):
+    """Raised when a safe insert or update fails due to a duplicate key error.
+    """
+
+
 class InvalidOperation(Exception):
     """Raised when a client attempts to perform an invalid operation.
     """

test/test_collection.py

@@ -31,6 +31,7 @@
 from pymongo.binary import Binary
 from pymongo.collection import Collection
 from pymongo.errors import InvalidName, OperationFailure, InvalidDocument
+from pymongo.errors import DuplicateKeyError
 from pymongo import ASCENDING, DESCENDING
 from pymongo.son import SON
 
@@ -462,6 +463,26 @@ def test_unique_index(self):
         db.test.save({"hello": "world"})
         self.assert_(db.error())
 
+    def test_duplicate_key_error(self):
+        db = self.db
+        db.drop_collection("test")
+
+        db.test.create_index("x", unique=True)
+
+        db.test.insert({"_id": 1, "x": 1}, safe=True)
+        db.test.insert({"_id": 2, "x": 2}, safe=True)
+
+        self.assertRaises(DuplicateKeyError,
+                          db.test.insert, {"_id": 1}, safe=True)
+        self.assertRaises(DuplicateKeyError,
+                          db.test.insert, {"x": 1}, safe=True)
+
+        self.assertRaises(DuplicateKeyError,
+                          db.test.save, {"x": 2}, safe=True)
+        self.assertRaises(DuplicateKeyError,
+                          db.test.update, {"x": 1}, {"$inc": {"x": 1}}, safe=True)
+
+
     def test_index_on_subfield(self):
         db = self.db
         db.drop_collection("test")