Merge pull request #967 from RussellLuo/master

Override `authentication_source` by "authSource" in URI

docs/changelog.rst

@@ -5,6 +5,7 @@ Changelog
 
 Changes in 0.9.X - DEV
 ======================
+- Added support for specifying authentication source as option `authSource` in URI. #967
 - Fixed mark_as_changed to handle higher/lower level fields changed. #927
 - ListField of embedded docs doesn't set the _instance attribute when iterating over it #914
 - Support += and *= for ListField #595

mongoengine/connection.py

@@ -1,4 +1,3 @@
-import pymongo
 from pymongo import MongoClient, MongoReplicaSetClient, uri_parser
 
 
@@ -58,8 +57,11 @@ def register_connection(alias, name=None, host=None, port=None,
             'password': uri_dict.get('password'),
             'read_preference': read_preference,
         })
-        if "replicaSet" in conn_settings['host']:
+        uri_options = uri_dict['options']
+        if 'replicaset' in uri_options:
             conn_settings['replicaSet'] = True
+        if 'authsource' in uri_options:
+            conn_settings['authentication_source'] = uri_options['authsource']
 
     # Deprecated parameters that should not be passed on
     kwargs.pop('slaves', None)

tests/test_connection.py

@@ -11,7 +11,10 @@ import datetime
 import pymongo
 from bson.tz_util import utc
 
-from mongoengine import *
+from mongoengine import (
+    connect, register_connection,
+    Document, DateTimeField
+)
 import mongoengine.connection
 from mongoengine.connection import get_db, get_connection, ConnectionError
 
@@ -101,6 +104,35 @@ class ConnectionTest(unittest.TestCase):
         c.admin.system.users.remove({})
         c.mongoenginetest.system.users.remove({})
 
+    def test_connect_uri_with_authsource(self):
+        """Ensure that the connect() method works well with
+        the option `authSource` in URI.
+        """
+        # Create users
+        c = connect('mongoenginetest')
+        c.admin.system.users.remove({})
+        c.admin.add_user('username', 'password')
+
+        # Authentication fails without "authSource"
+        self.assertRaises(
+            ConnectionError, connect, 'mongoenginetest', alias='test1',
+            host='mongodb://username:password@localhost/mongoenginetest'
+        )
+        self.assertRaises(ConnectionError, get_db, 'test1')
+
+        # Authentication succeeds with "authSource"
+        connect(
+            'mongoenginetest', alias='test2',
+            host=('mongodb://username:password@localhost/'
+                  'mongoenginetest?authSource=admin')
+        )
+        db = get_db('test2')
+        self.assertTrue(isinstance(db, pymongo.database.Database))
+        self.assertEqual(db.name, 'mongoenginetest')
+
+        # Clear all users
+        c.admin.system.users.remove({})
+
     def test_register_connection(self):
         """Ensure that connections with different aliases may be registered.
         """