diff --git a/mongoengine/connection.py b/mongoengine/connection.py index 34ff4dc31..dd35fd2dc 100644 --- a/mongoengine/connection.py +++ b/mongoengine/connection.py @@ -92,17 +92,21 @@ def register_connection(alias, name=None, host=None, port=None, if uri_dict.get('database'): conn_settings['name'] = uri_dict.get('database') - for param in ('read_preference', 'username', 'password'): + for param in ('username', 'password'): if uri_dict.get(param): conn_settings[param] = uri_dict[param] uri_options = uri_dict['options'] - if 'replicaset' in uri_options: - conn_settings['replicaSet'] = uri_options['replicaset'] + if 'replicaset' in uri_options \ + or 'replicaset' in [kw_key.lower() for kw_key in kwargs.keys()]: + conn_settings['isreplicaSet'] = True if 'authsource' in uri_options: conn_settings['authentication_source'] = uri_options['authsource'] if 'authmechanism' in uri_options: conn_settings['authentication_mechanism'] = uri_options['authmechanism'] + if IS_PYMONGO_3 and 'readpreference' in uri_options \ + and 'readpreferencetags' in uri_options: + del conn_settings['read_preference'] else: resolved_hosts.append(entity) conn_settings['host'] = resolved_hosts @@ -176,19 +180,18 @@ def _clean_settings(settings_dict): # For replica set connections with PyMongo 2.x, use # MongoReplicaSetClient. # TODO remove this once we stop supporting PyMongo 2.x. - if 'replicaSet' in conn_settings and not IS_PYMONGO_3: - connection_class = MongoReplicaSetClient - conn_settings['hosts_or_uri'] = conn_settings.pop('host', None) - - # hosts_or_uri has to be a string, so if 'host' was provided - # as a list, join its parts and separate them by ',' - if isinstance(conn_settings['hosts_or_uri'], list): - conn_settings['hosts_or_uri'] = ','.join( - conn_settings['hosts_or_uri']) - - # Discard port since it can't be used on MongoReplicaSetClient - conn_settings.pop('port', None) - + if 'isreplicaSet' in conn_settings.keys(): + if not IS_PYMONGO_3: + connection_class = MongoReplicaSetClient + conn_settings['hosts_or_uri'] = conn_settings.pop('host', None) + # hosts_or_uri has to be a string, so if 'host' was provided + # as a list, join its parts and separate them by ',' + if isinstance(conn_settings['hosts_or_uri'], list): + conn_settings['hosts_or_uri'] = ','.join( + conn_settings['hosts_or_uri']) + # Discard port since it can't be used on MongoReplicaSetClient + conn_settings.pop('port', None) + del conn_settings['isreplicaSet'] # Iterate over all of the connection settings and if a connection with # the same parameters is already established, use it instead of creating # a new one. diff --git a/tests/connections/__init__.py b/tests/connections/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/test_connection.py b/tests/connections/test_connection.py similarity index 82% rename from tests/test_connection.py rename to tests/connections/test_connection.py index cdcf13773..b8cc972d6 100644 --- a/tests/test_connection.py +++ b/tests/connections/test_connection.py @@ -211,6 +211,78 @@ def test_uri_without_credentials_doesnt_override_conn_settings(self): # behavior. If the MongoDB URI would override the credentials self.assertRaises(OperationFailure, get_db) + @unittest.skip('require authentication enabled and meet needed prerequisites') + def test_uri_overide_kwargs_with_right_authSource_not_pass(self): + """Authenticated during building connection""" + ###as prerequisite, require two users at the time when authentication is disabled### + #no_auth_conn_obj = connect(db='mongoenginetest') + #no_auth_conn_obj.admin.system.users.remove({}) + #no_auth_conn_obj.admin.add_user('rjxu', 'rjxp', read_only=True) + #no_auth_conn_obj.mongoenginetest.add_user('rjxtu','rjxtp', read_only=True) + + auth_conn_obj = connect( + db='mongoenginetest', + alias='authtest', + username='rjxu', + password='rjxp', + host="mongodb://rjxtu:rjxtp@localhost:27017" + + "/?authSource=mongoenginetest", + authentication_source='admin', + authentication_mechanism='MONGODB-X509' + ) + + if not IS_PYMONGO_3: + #attempt to read data from a collection authored by admin + def get_data_case(): + auth_db_obj = auth_conn_obj['admin'] + coll_from_auth_db = auth_db_obj['system.users'] + return [sys_user for sys_user in coll_from_auth_db.find()] + try: + get_data_case() + except: + self.assertRaises(OperationFailure, get_data_case) + else: + credentials_obj = auth_conn_obj._MongoClient__options.credentials + self.assertEqual(credentials_obj.username, 'rjxtu') + self.assertEqual(credentials_obj.source, 'mongoenginetest') + self.assertEqual(credentials_obj.mechanism, 'DEFAULT') + + @unittest.skip('require authentication enabled and meet needed prerequisite') + def test_post_auth_mongo_client_with_get_db(self): + '''Post authentication at get db stage''' + ###as prerequisite, require one user at the time when authentication is disabled### + #no_auth_conn_obj = connect(db='mongoenginetest') + #no_auth_conn_obj.admin.system.users.remove({}) + #no_auth_conn_obj.mongoenginetest.add_user('rjxtu', 'rjxtp', read_only=True) + + no_auth_conn_obj = connect( + db='mongoenginetest', + username='rjxtu', + password='rjxtp', + host="mongodb://localhost:27017", + authentication_source='mongoenginetest', + authentication_mechanism='DEFAULT' + ) + + #read data from auth or non-auth conn + def get_data_case(is_get_db=False): + if not is_get_db: + _db_obj = no_auth_conn_obj['mongoenginetest'] + else: + _db_obj = get_db() + coll_from_db = _db_obj['testcoll'] + return [_td for _td in coll_from_db.find()] + + # + try: + get_data_case() + except: + self.assertRaises(OperationFailure, get_data_case) + # + self.assertRaises(AssertionError, self.assertRaises, + OperationFailure, get_data_case, True) + + def test_connect_uri_with_authsource(self): """Ensure that the connect() method works well with `authSource` option in the URI. @@ -319,38 +391,6 @@ def test_write_concern(self): self.assertEqual(dict(conn1.write_concern), {'w': 1, 'j': True}) self.assertEqual(dict(conn2.write_concern), {'w': 1, 'j': True}) - def test_connect_with_replicaset_via_uri(self): - """Ensure connect() works when specifying a replicaSet via the - MongoDB URI. - """ - if IS_PYMONGO_3: - c = connect(host='mongodb://localhost/test?replicaSet=local-rs') - db = get_db() - self.assertTrue(isinstance(db, pymongo.database.Database)) - self.assertEqual(db.name, 'test') - else: - # PyMongo < v3.x raises an exception: - # "localhost:27017 is not a member of replica set local-rs" - with self.assertRaises(MongoEngineConnectionError): - c = connect(host='mongodb://localhost/test?replicaSet=local-rs') - - def test_connect_with_replicaset_via_kwargs(self): - """Ensure connect() works when specifying a replicaSet via the - connection kwargs - """ - if IS_PYMONGO_3: - c = connect(replicaset='local-rs') - self.assertEqual(c._MongoClient__options.replica_set_name, - 'local-rs') - db = get_db() - self.assertTrue(isinstance(db, pymongo.database.Database)) - self.assertEqual(db.name, 'test') - else: - # PyMongo < v3.x raises an exception: - # "localhost:27017 is not a member of replica set local-rs" - with self.assertRaises(MongoEngineConnectionError): - c = connect(replicaset='local-rs') - def test_datetime(self): connect('mongoenginetest', tz_aware=True) d = datetime.datetime(2010, 5, 5, tzinfo=utc) diff --git a/tests/connections/test_replicaset_connection.py b/tests/connections/test_replicaset_connection.py new file mode 100644 index 000000000..62e707e2e --- /dev/null +++ b/tests/connections/test_replicaset_connection.py @@ -0,0 +1,147 @@ +import unittest + +from pymongo import (ReadPreference, + read_preferences) + +from mongoengine.python_support import IS_PYMONGO_3 + +if IS_PYMONGO_3: + from pymongo import MongoClient + CONN_CLASS = MongoClient + READ_PREF = ReadPreference.SECONDARY +else: + from pymongo import ReplicaSetConnection + CONN_CLASS = ReplicaSetConnection + READ_PREF = ReadPreference.SECONDARY_ONLY + +import pymongo +import mongoengine +from mongoengine import * +from mongoengine.connection import (MongoEngineConnectionError, + get_db) + + +class ConnectionTest(unittest.TestCase): + + def setUp(self): + mongoengine.connection._connection_settings = {} + mongoengine.connection._connections = {} + mongoengine.connection._dbs = {} + + def tearDown(self): + mongoengine.connection._connection_settings = {} + mongoengine.connection._connections = {} + mongoengine.connection._dbs = {} + + def test_replicaset_uri_passes_read_preference(self): + """Requires a replica set called "rs" on port 27017 + """ + + try: + conn = connect(db='mongoenginetest', + host="mongodb://localhost/mongoenginetest?replicaSet=local-rs", + read_preference=READ_PREF) + except MongoEngineConnectionError as e: + return + + if not isinstance(conn, CONN_CLASS): + # really??? + return + + self.assertEqual(conn.read_preference, READ_PREF) + + def test_connect_with_replicaset_via_kwargs(self): + """Ensure connect() works when specifying a replicaSet via the + connection kwargs + """ + if IS_PYMONGO_3: + c = connect(replicaset='local-rs') + self.assertEqual(c._MongoClient__options.replica_set_name, + 'local-rs') + db = get_db() + self.assertTrue(isinstance(db, pymongo.database.Database)) + self.assertEqual(db.name, 'test') + else: + # PyMongo < v3.x raises an exception: + # "localhost:27017 is not a member of replica set local-rs" + with self.assertRaises(MongoEngineConnectionError): + c = connect(replicaset='local-rs') + + def test_connect_with_replicaset_via_uri(self): + """Ensure connect() works when specifying a replicaSet via the + MongoDB URI. + """ + if IS_PYMONGO_3: + c = connect(host='mongodb://localhost/test?replicaSet=local-rs') + db = get_db() + self.assertTrue(isinstance(db, pymongo.database.Database)) + self.assertEqual(db.name, 'test') + else: + # PyMongo < v3.x raises an exception: + # "localhost:27017 is not a member of replica set local-rs" + with self.assertRaises(MongoEngineConnectionError): + c = connect(host='mongodb://localhost/test?replicaSet=local-rs') + + def test_read_preference_from_replica_set_in_uri_as_host(self): + '''read preference from replica set cluster''' + #test case about uri option overrides kwargs and read_preference with tag sets + if IS_PYMONGO_3: + conn_obj = connect( + db='testrjx', + host="mongodb://localhost:27017/?replicaset=dev_rs" + + "&readpreference=nearest&readpreferencetags=", + read_preference=ReadPreference.SECONDARY + ) + read_preference_obj = conn_obj.read_preference + self.assertEqual(read_preference_obj.mode, ReadPreference.NEAREST.mode) + self.assertEqual(read_preference_obj.tag_sets, [{}]) + else: + if pymongo.version_tuple[1] < 9: # like v2.8 + with self.assertRaises(MongoEngineConnectionError): + c_obj = connect( + db='testrjx', + host="mongodb://localhost:27018/?replicaset=dev_rs" + + "&read_preference=nearest&readpreferencetags=dc:east,use:dev", + read_preference=ReadPreference.SECONDARY, + readpreferencetags=[{}] #tag_sets as an alternative + ) + self.assertEqual(c_obj.read_preference, ReadPreference.NEAREST) + self.assertEqual(c_obj.tag_sets, [{'dc':'east','use':'dev'}]) + else: + #for 2.9 as transition version, also accept read_preference object subclassing _ServerMode + with self.assertRaises(MongoEngineConnectionError): + c_obj = connect( + db='testrjx', + host="mongodb://localhost:27018/?replicaset=dev_rs", + read_preference=read_preferences.SecondaryPreferred( + [{'dc':'east','use':'dev'}] + ), + ) + self.assertEqual(c_obj.read_preference, ReadPreference.SECONDARY_PREFERRED) + self.assertEqual(c_obj.tag_sets, [{'dc':'east','use':'dev'}]) + + def test_replica_set_as_uri_or_kwargs(self): + '''replicaset as uri option or kwargs''' + # dev_rs as real replica set name, while test_rs is false as comparision + # replicaset case insensitive + if not IS_PYMONGO_3: #like 2.9 and 2.8 + # uri opt overrides kwargs for connecting replica set + with self.assertRaises(MongoEngineConnectionError): + c_obj = connect( + db='testrjx', + host="mongodb://localhost:27018/?replicaset=dev_rs", + replicaset='test_rs' + ) + self.assertEqual(c_obj._MongoReplicaSetClient__name, 'dev_rs') + else: #like 3.3 + #here, kwargs overrides uri option + #though wrong replica set to be connected, it is ok until issuing real data operation + c_obj = connect( + db='testrjx', + host="mongodb://localhost:27017/?replicaSet=dev_rs", + replicaSet='test_rs' + ) + self.assertEqual(c_obj._MongoClient__options.replica_set_name, 'test_rs') + +if __name__ == '__main__': + unittest.main() diff --git a/tests/test_replicaset_connection.py b/tests/test_replicaset_connection.py deleted file mode 100644 index a53f5903e..000000000 --- a/tests/test_replicaset_connection.py +++ /dev/null @@ -1,51 +0,0 @@ -import unittest - -from pymongo import ReadPreference - -from mongoengine.python_support import IS_PYMONGO_3 - -if IS_PYMONGO_3: - from pymongo import MongoClient - CONN_CLASS = MongoClient - READ_PREF = ReadPreference.SECONDARY -else: - from pymongo import ReplicaSetConnection - CONN_CLASS = ReplicaSetConnection - READ_PREF = ReadPreference.SECONDARY_ONLY - -import mongoengine -from mongoengine import * -from mongoengine.connection import MongoEngineConnectionError - - -class ConnectionTest(unittest.TestCase): - - def setUp(self): - mongoengine.connection._connection_settings = {} - mongoengine.connection._connections = {} - mongoengine.connection._dbs = {} - - def tearDown(self): - mongoengine.connection._connection_settings = {} - mongoengine.connection._connections = {} - mongoengine.connection._dbs = {} - - def test_replicaset_uri_passes_read_preference(self): - """Requires a replica set called "rs" on port 27017 - """ - - try: - conn = connect(db='mongoenginetest', - host="mongodb://localhost/mongoenginetest?replicaSet=rs", - read_preference=READ_PREF) - except MongoEngineConnectionError as e: - return - - if not isinstance(conn, CONN_CLASS): - # really??? - return - - self.assertEqual(conn.read_preference, READ_PREF) - -if __name__ == '__main__': - unittest.main()