Merge pull request #2068 from bagerard/fix_connection_auth_same_host

Fix connection issue when using different authentication in different dbs
This commit is contained in:
Bastien Gérard 2019-06-07 21:08:26 +02:00 committed by GitHub
commit 0ac59c67ea
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 59 additions and 22 deletions

View File

@ -13,6 +13,7 @@ Development
- expose `mongoengine.connection.disconnect` and `mongoengine.connection.disconnect_all`
- Fix disconnect function #566 #1599 #605 #607 #1213 #565
- Improve connect/disconnect documentations
- Fix issue when using multiple connections to the same mongo with different credentials #2047
- POTENTIAL BREAKING CHANGES: (associated with connect/disconnect fixes)
- calling `connect` 2 times with the same alias and different parameter will raise an error (should call disconnect first)
- disconnect now clears `mongoengine.connection._connection_settings`

View File

@ -235,7 +235,6 @@ def get_connection(alias=DEFAULT_CONNECTION_NAME, reconnect=False):
raise MongoEngineConnectionError(msg)
def _clean_settings(settings_dict):
# set literal more efficient than calling set function
irrelevant_fields_set = {
'name', 'username', 'password',
'authentication_source', 'authentication_mechanism'
@ -245,10 +244,12 @@ def get_connection(alias=DEFAULT_CONNECTION_NAME, reconnect=False):
if k not in irrelevant_fields_set
}
raw_conn_settings = _connection_settings[alias].copy()
# Retrieve a copy of the connection settings associated with the requested
# alias and remove the database name and authentication info (we don't
# care about them at this point).
conn_settings = _clean_settings(_connection_settings[alias].copy())
conn_settings = _clean_settings(raw_conn_settings)
# Determine if we should use PyMongo's or mongomock's MongoClient.
is_mock = conn_settings.pop('is_mock', False)
@ -262,33 +263,58 @@ def get_connection(alias=DEFAULT_CONNECTION_NAME, reconnect=False):
else:
connection_class = MongoClient
# 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.
existing_connection = None
connection_settings_iterator = (
(db_alias, settings.copy())
for db_alias, settings in _connection_settings.items()
)
for db_alias, connection_settings in connection_settings_iterator:
connection_settings = _clean_settings(connection_settings)
if conn_settings == connection_settings and _connections.get(db_alias):
existing_connection = _connections[db_alias]
break
# Re-use existing connection if one is suitable
existing_connection = _find_existing_connection(raw_conn_settings)
# If an existing connection was found, assign it to the new alias
if existing_connection:
_connections[alias] = existing_connection
else:
# Otherwise, create the new connection for this alias. Raise
# MongoEngineConnectionError if it can't be established.
_connections[alias] = _create_connection(alias=alias,
connection_class=connection_class,
**conn_settings)
return _connections[alias]
def _create_connection(alias, connection_class, **connection_settings):
"""
Create the new connection for this alias. Raise
MongoEngineConnectionError if it can't be established.
"""
try:
_connections[alias] = connection_class(**conn_settings)
return connection_class(**connection_settings)
except Exception as e:
raise MongoEngineConnectionError(
'Cannot connect to database %s :\n%s' % (alias, e))
return _connections[alias]
def _find_existing_connection(connection_settings):
"""
Check if an existing connection could be reused
Iterate over all of the connection settings and if an existing connection
with the same parameters is suitable, return it
:param connection_settings: the settings of the new connection
:return: An existing connection or None
"""
connection_settings_bis = (
(db_alias, settings.copy())
for db_alias, settings in _connection_settings.items()
)
def _clean_settings(settings_dict):
# Only remove the name but it's important to
# keep the username/password/authentication_source/authentication_mechanism
# to identify if the connection could be shared (cfr https://github.com/MongoEngine/mongoengine/issues/2047)
return {k: v for k, v in settings_dict.items() if k != 'name'}
cleaned_conn_settings = _clean_settings(connection_settings)
for db_alias, connection_settings in connection_settings_bis:
db_conn_settings = _clean_settings(connection_settings)
if cleaned_conn_settings == db_conn_settings and _connections.get(db_alias):
return _connections[db_alias]
def get_db(alias=DEFAULT_CONNECTION_NAME, reconnect=False):

View File

@ -611,6 +611,16 @@ class ConnectionTest(unittest.TestCase):
self.assertEqual(mongo_connections['t1'].address[0], 'localhost')
self.assertEqual(mongo_connections['t2'].address[0], '127.0.0.1')
def test_connect_2_databases_uses_same_client_if_only_dbname_differs(self):
c1 = connect(alias='testdb1', db='testdb1')
c2 = connect(alias='testdb2', db='testdb2')
self.assertIs(c1, c2)
def test_connect_2_databases_uses_different_client_if_different_parameters(self):
c1 = connect(alias='testdb1', db='testdb1', username='u1')
c2 = connect(alias='testdb2', db='testdb2', username='u2')
self.assertIsNot(c1, c2)
if __name__ == '__main__':
unittest.main()