From 1c5e6a3425f65f63c36da62a2a3aa0f0584cbfd4 Mon Sep 17 00:00:00 2001
From: Ross Lawley <ross.lawley@gmail.com>
Date: Fri, 24 Aug 2012 10:38:00 +0100
Subject: [PATCH] NotUniqueError gracefully replacing ambiguous OperationError
 when appropriate

---
 mongoengine/document.py | 10 +++++++---
 mongoengine/queryset.py |  9 ++++++++-
 tests/test_document.py  | 11 +++++++----
 tests/test_queryset.py  |  2 +-
 4 files changed, 23 insertions(+), 9 deletions(-)

diff --git a/mongoengine/document.py b/mongoengine/document.py
index 4fbf1fe0..1691df93 100644
--- a/mongoengine/document.py
+++ b/mongoengine/document.py
@@ -1,18 +1,19 @@
 import warnings
 
 import pymongo
+import re
 
 from bson.dbref import DBRef
 from mongoengine import signals, queryset
 
 from base import (DocumentMetaclass, TopLevelDocumentMetaclass, BaseDocument,
                   BaseDict, BaseList)
-from queryset import OperationError
+from queryset import OperationError, NotUniqueError
 from connection import get_db, DEFAULT_CONNECTION_NAME
 
 __all__ = ['Document', 'EmbeddedDocument', 'DynamicDocument',
            'DynamicEmbeddedDocument', 'OperationError',
-           'InvalidCollectionError']
+           'InvalidCollectionError', 'NotUniqueError']
 
 
 class InvalidCollectionError(Exception):
@@ -250,8 +251,11 @@ class Document(BaseDocument):
 
         except pymongo.errors.OperationFailure, err:
             message = 'Could not save document (%s)'
-            if u'duplicate key' in unicode(err):
+            if re.match('^E1100[01] duplicate key', unicode(err)):
+                # E11000 - duplicate key error index
+                # E11001 - duplicate key on update
                 message = u'Tried to save duplicate unique keys (%s)'
+                raise NotUniqueError(message % unicode(err))
             raise OperationError(message % unicode(err))
         id_field = self._meta['id_field']
         if id_field not in self._meta.get('shard_key', []):
diff --git a/mongoengine/queryset.py b/mongoengine/queryset.py
index b9106c2b..ecd58936 100644
--- a/mongoengine/queryset.py
+++ b/mongoengine/queryset.py
@@ -45,6 +45,10 @@ class OperationError(Exception):
     pass
 
 
+class NotUniqueError(OperationError):
+    pass
+
+
 RE_TYPE = type(re.compile(''))
 
 
@@ -924,8 +928,11 @@ class QuerySet(object):
             ids = self._collection.insert(raw, **write_options)
         except pymongo.errors.OperationFailure, err:
             message = 'Could not save document (%s)'
-            if u'duplicate key' in unicode(err):
+            if re.match('^E1100[01] duplicate key', unicode(err)):
+                # E11000 - duplicate key error index
+                # E11001 - duplicate key on update
                 message = u'Tried to save duplicate unique keys (%s)'
+                raise NotUniqueError(message % unicode(err))
             raise OperationError(message % unicode(err))
 
         if not load_bulk:
diff --git a/tests/test_document.py b/tests/test_document.py
index 274bbb99..5faf2680 100644
--- a/tests/test_document.py
+++ b/tests/test_document.py
@@ -1013,6 +1013,9 @@ class DocumentTest(unittest.TestCase):
 
         # Two posts with the same slug is not allowed
         post2 = BlogPost(title='test2', slug='test')
+        self.assertRaises(NotUniqueError, post2.save)
+
+        # Ensure backwards compatibilty for errors
         self.assertRaises(OperationError, post2.save)
 
     def test_unique_with(self):
@@ -1063,7 +1066,7 @@ class DocumentTest(unittest.TestCase):
 
         # Now there will be two docs with the same sub.slug
         post3 = BlogPost(title='test3', sub=SubDocument(year=2010, slug='test'))
-        self.assertRaises(OperationError, post3.save)
+        self.assertRaises(NotUniqueError, post3.save)
 
         BlogPost.drop_collection()
 
@@ -1090,11 +1093,11 @@ class DocumentTest(unittest.TestCase):
 
         # Now there will be two docs with the same sub.slug
         post3 = BlogPost(title='test3', sub=SubDocument(year=2010, slug='test'))
-        self.assertRaises(OperationError, post3.save)
+        self.assertRaises(NotUniqueError, post3.save)
 
         # Now there will be two docs with the same title and year
         post3 = BlogPost(title='test1', sub=SubDocument(year=2009, slug='test-1'))
-        self.assertRaises(OperationError, post3.save)
+        self.assertRaises(NotUniqueError, post3.save)
 
         BlogPost.drop_collection()
 
@@ -1117,7 +1120,7 @@ class DocumentTest(unittest.TestCase):
         try:
             cust_dupe.save()
             raise AssertionError, "We saved a dupe!"
-        except OperationError:
+        except NotUniqueError:
             pass
         Customer.drop_collection()
 
diff --git a/tests/test_queryset.py b/tests/test_queryset.py
index 8c998e5e..8f0ad94e 100644
--- a/tests/test_queryset.py
+++ b/tests/test_queryset.py
@@ -578,7 +578,7 @@ class QuerySetTest(unittest.TestCase):
         def throw_operation_error_not_unique():
             Blog.objects.insert([blog2, blog3], safe=True)
 
-        self.assertRaises(OperationError, throw_operation_error_not_unique)
+        self.assertRaises(NotUniqueError, throw_operation_error_not_unique)
         self.assertEqual(Blog.objects.count(), 2)
 
         Blog.objects.insert([blog2, blog3], write_options={'continue_on_error': True})