From 9634e44343b31e0880aeb628ac5e89dba3a4d97a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Bastien=20G=C3=A9rard?= Date: Wed, 15 May 2019 21:43:29 +0200 Subject: [PATCH 1/3] Fix the issue that the same MongoClient gets re-used in case we connect to 2 databases on the same host (problematic when different users authenticate) --- mongoengine/connection.py | 57 ++++++++++++++++++++++++++++----------- 1 file changed, 42 insertions(+), 15 deletions(-) diff --git a/mongoengine/connection.py b/mongoengine/connection.py index e0399fde..5fae9507 100644 --- a/mongoengine/connection.py +++ b/mongoengine/connection.py @@ -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,11 @@ 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,19 +262,8 @@ 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: @@ -291,6 +280,44 @@ def get_connection(alias=DEFAULT_CONNECTION_NAME, reconnect=False): return _connections[alias] +def _create_connection(connection_class, **connection_settings): + # Otherwise, create the new connection for this alias. Raise + # MongoEngineConnectionError if it can't be established. + try: + _connections[alias] = connection_class(**conn_settings) + except Exception as e: + raise MongoEngineConnectionError( + 'Cannot connect to database %s :\n%s' % (alias, e)) + + +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_iterator = ( + (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_iterator: + 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): if reconnect: disconnect(alias) From 84c42ed58c6c7d26ec6a392048653d1993acb319 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Bastien=20G=C3=A9rard?= Date: Tue, 4 Jun 2019 22:35:42 +0200 Subject: [PATCH 2/3] Add tests --- mongoengine/connection.py | 25 ++++++++++++------------- tests/test_connection.py | 10 ++++++++++ 2 files changed, 22 insertions(+), 13 deletions(-) diff --git a/mongoengine/connection.py b/mongoengine/connection.py index 5fae9507..6249225c 100644 --- a/mongoengine/connection.py +++ b/mongoengine/connection.py @@ -245,6 +245,7 @@ def get_connection(alias=DEFAULT_CONNECTION_NAME, reconnect=False): } 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). @@ -269,22 +270,20 @@ def get_connection(alias=DEFAULT_CONNECTION_NAME, reconnect=False): if existing_connection: _connections[alias] = existing_connection else: - # Otherwise, create the new connection for this alias. Raise - # MongoEngineConnectionError if it can't be established. - try: - _connections[alias] = connection_class(**conn_settings) - except Exception as e: - raise MongoEngineConnectionError( - 'Cannot connect to database %s :\n%s' % (alias, e)) + _create_connection(alias=alias, + connection_class=connection_class, + **conn_settings) return _connections[alias] -def _create_connection(connection_class, **connection_settings): - # Otherwise, create the new connection for this alias. Raise - # MongoEngineConnectionError if it can't be established. +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) + _connections[alias] = connection_class(**connection_settings) except Exception as e: raise MongoEngineConnectionError( 'Cannot connect to database %s :\n%s' % (alias, e)) @@ -300,7 +299,7 @@ def _find_existing_connection(connection_settings): :param connection_settings: the settings of the new connection :return: An existing connection or None """ - connection_settings_iterator = ( + connection_settings_bis = ( (db_alias, settings.copy()) for db_alias, settings in _connection_settings.items() ) @@ -312,7 +311,7 @@ def _find_existing_connection(connection_settings): 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_iterator: + 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] diff --git a/tests/test_connection.py b/tests/test_connection.py index 5473b8a0..d3fcc395 100644 --- a/tests/test_connection.py +++ b/tests/test_connection.py @@ -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() From 36aebffcc0f18fd6823d734cdd4001cef7474f9d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Bastien=20G=C3=A9rard?= Date: Tue, 4 Jun 2019 22:39:44 +0200 Subject: [PATCH 3/3] update changelog --- docs/changelog.rst | 1 + mongoengine/connection.py | 8 ++++---- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/docs/changelog.rst b/docs/changelog.rst index 87729df3..7e8fd3d2 100644 --- a/docs/changelog.rst +++ b/docs/changelog.rst @@ -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` diff --git a/mongoengine/connection.py b/mongoengine/connection.py index 6249225c..9d4f25fc 100644 --- a/mongoengine/connection.py +++ b/mongoengine/connection.py @@ -270,9 +270,9 @@ def get_connection(alias=DEFAULT_CONNECTION_NAME, reconnect=False): if existing_connection: _connections[alias] = existing_connection else: - _create_connection(alias=alias, - connection_class=connection_class, - **conn_settings) + _connections[alias] = _create_connection(alias=alias, + connection_class=connection_class, + **conn_settings) return _connections[alias] @@ -283,7 +283,7 @@ def _create_connection(alias, connection_class, **connection_settings): MongoEngineConnectionError if it can't be established. """ try: - _connections[alias] = connection_class(**connection_settings) + return connection_class(**connection_settings) except Exception as e: raise MongoEngineConnectionError( 'Cannot connect to database %s :\n%s' % (alias, e))