Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions docs/changelog.rst
Original file line number Diff line number Diff line change
Expand Up @@ -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`
Expand Down
70 changes: 48 additions & 22 deletions mongoengine/connection.py
Original file line number Diff line number Diff line change
Expand Up @@ -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'
Expand All @@ -245,10 +244,12 @@ def _clean_settings(settings_dict):
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)
Expand All @@ -262,35 +263,60 @@ def _clean_settings(settings_dict):
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.
try:
_connections[alias] = connection_class(**conn_settings)
except Exception as e:
raise MongoEngineConnectionError(
'Cannot connect to database %s :\n%s' % (alias, e))
_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:
return connection_class(**connection_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_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):
if reconnect:
disconnect(alias)
Expand Down
10 changes: 10 additions & 0 deletions tests/test_connection.py
Original file line number Diff line number Diff line change
Expand Up @@ -611,6 +611,16 @@ def test_multiple_connection_settings(self):
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()