Multiposting changes to MongoEngine

docs/guide/connecting.rst

@@ -23,6 +23,11 @@ arguments should be provided::
 
     connect('project1', username='webapp', password='pwd123')
 
+The :attr:`authsource` arguments allows you to specify the database to
+authentication. If not specified the current database is used::
+
+    connect('project1', username='webapp', password='pwd123', authsource='admin')
+
 Uri style connections are also supported - just supply the uri as
 the :attr:`host` to
 :func:`~mongoengine.connect`::

mongoengine/connection.py

@@ -20,7 +20,8 @@ class ConnectionError(Exception):
 
 def register_connection(alias, name, host=None, port=None,
                         is_slave=False, read_preference=False, slaves=None,
-                        username=None, password=None, **kwargs):
+                        username=None, password=None, authsource=None,
+                        **kwargs):
     """Add a connection.
 
     :param alias: the name that will be used to refer to this connection
@@ -36,6 +37,8 @@ def register_connection(alias, name, host=None, port=None,
         be a registered connection that has :attr:`is_slave` set to ``True``
     :param username: username to authenticate with
     :param password: password to authenticate with
+    :param authsource: database to authenticate on
+       ** Added pymongo 2.5
     :param kwargs: allow ad-hoc parameters to be passed into the pymongo driver
 
     """
@@ -49,6 +52,7 @@ def register_connection(alias, name, host=None, port=None,
         'slaves': slaves or [],
         'username': username,
         'password': password,
+        'authsource': authsource,
         'read_preference': read_preference
     }
 
@@ -99,6 +103,7 @@ def get_connection(alias=DEFAULT_CONNECTION_NAME, reconnect=False):
             conn_settings.pop('is_slave', None)
             conn_settings.pop('username', None)
             conn_settings.pop('password', None)
+            conn_settings.pop('authsource', None)
         else:
             # Get all the slave connections
             if 'slaves' in conn_settings:
@@ -137,7 +142,8 @@ def get_db(alias=DEFAULT_CONNECTION_NAME, reconnect=False):
         # Authenticate if necessary
         if conn_settings['username'] and conn_settings['password']:
             db.authenticate(conn_settings['username'],
-                            conn_settings['password'])
+                            conn_settings['password'],
+                            conn_settings['authsource'])
         _dbs[alias] = db
     return _dbs[alias]
 

mongoengine/document.py

@@ -280,6 +280,9 @@ def is_new_object(last_error):
                     kwargs.update(cascade_kwargs)
                 kwargs['_refs'] = _refs
                 self.cascade_save(**kwargs)
+        except pymongo.errors.DuplicateKeyError, err:
+            message = u'Tried to save duplicate unique keys (%s)'
+            raise NotUniqueError(message % unicode(err))
 
         except pymongo.errors.OperationFailure, err:
             message = 'Could not save document (%s)'

mongoengine/queryset/base.py

@@ -302,6 +302,9 @@ def insert(self, doc_or_docs, load_bulk=True, write_concern=None):
         signals.pre_bulk_insert.send(self._document, documents=docs)
         try:
             ids = self._collection.insert(raw, **write_concern)
+        except pymongo.errors.DuplicateKeyError, err:
+            message = 'Could not save document (%s)'
+            raise NotUniqueError(message % unicode(err))
         except pymongo.errors.OperationFailure, err:
             message = 'Could not save document (%s)'
             if re.match('^E1100[01] duplicate key', unicode(err)):
@@ -440,6 +443,8 @@ def update(self, upsert=False, multi=True, write_concern=None,
                 return result
             elif result:
                 return result['n']
+        except pymongo.errors.DuplicateKeyError, err:
+            raise NotUniqueError(u'Update failed (%s)' % unicode(err))
         except pymongo.errors.OperationFailure, err:
             if unicode(err) == u'multi not coded yet':
                 message = u'update() method requires MongoDB 1.1.3+'

tests/document/instance.py

@@ -15,7 +15,7 @@
 
 from mongoengine import *
 from mongoengine.errors import (NotRegistered, InvalidDocumentError,
-                                InvalidQueryError)
+                                InvalidQueryError, NotUniqueError)
 from mongoengine.queryset import NULLIFY, Q
 from mongoengine.connection import get_db
 from mongoengine.base import get_document
@@ -969,6 +969,16 @@ def update_no_op_raises():
 
         self.assertRaises(InvalidQueryError, update_no_op_raises)
 
+    def test_update_unique_field(self):
+        class Doc(Document):
+            name = StringField(unique=True)
+
+        doc1 = Doc(name="first").save()
+        doc2 = Doc(name="second").save()
+
+        self.assertRaises(NotUniqueError, lambda:
+                          doc2.update(set__name=doc1.name))
+
     def test_embedded_update(self):
         """
         Test update on `EmbeddedDocumentField` fields