Merge pull request #2716 from bagerard/ShaneHarvey-fix-disconnect-closes-clients-too-early

Clone - Shane harvey fix disconnect closes clients too early

Author: bagerard

docs/changelog.rst

@@ -11,6 +11,8 @@ Development
 - Turning off dereferencing for the results of distinct query. #2663
 - Add tests against Mongo 5.0 in pipeline
 - Drop support for Python 3.6 (EOL)
+- Bug fix support for PyMongo>=4 to fix "InvalidOperation: Cannot use MongoClient after close"
+  errors.
 
 Changes in 0.24.2
 =================

mongoengine/connection.py

@@ -246,9 +246,15 @@ def disconnect(alias=DEFAULT_CONNECTION_NAME):
     from mongoengine import Document
     from mongoengine.base.common import _get_documents_by_db
 
-    if alias in _connections:
-        get_connection(alias=alias).close()
-        del _connections[alias]
+    connection = _connections.pop(alias, None)
+    if connection:
+        # MongoEngine may share the same MongoClient across multiple aliases
+        # if connection settings are the same so we only close
+        # the client if we're removing the final reference.
+        # Important to use 'is' instead of '==' because clients connected to the same cluster
+        # will compare equal even with different options
+        if all(connection is not c for c in _connections.values()):
+            connection.close()
 
     if alias in _dbs:
         # Detach all cached collections in Documents

tests/test_connection.py

@@ -5,7 +5,11 @@
 import pytest
 from bson.tz_util import utc
 from pymongo import MongoClient, ReadPreference
-from pymongo.errors import InvalidName, OperationFailure
+from pymongo.errors import (
+    InvalidName,
+    InvalidOperation,
+    OperationFailure,
+)
 
 import mongoengine.connection
 from mongoengine import (
@@ -287,6 +291,30 @@ def test_disconnect_silently_pass_if_alias_does_not_exist(self):
         assert len(connections) == 0
         disconnect(alias="not_exist")
 
+    def test_disconnect_does_not_close_client_used_by_another_alias(self):
+        client1 = connect(alias="disconnect_reused_client_test_1")
+        client2 = connect(alias="disconnect_reused_client_test_2")
+        client3 = connect(alias="disconnect_reused_client_test_3", maxPoolSize=10)
+        assert client1 is client2
+        assert client1 is not client3
+        client1.admin.command("ping")
+        disconnect("disconnect_reused_client_test_1")
+        # The client is not closed because the second alias still exists.
+        client2.admin.command("ping")
+        disconnect("disconnect_reused_client_test_2")
+        # The client is now closed:
+        if PYMONGO_VERSION >= (4,):
+            with pytest.raises(InvalidOperation):
+                client2.admin.command("ping")
+        # 3rd client connected to the same cluster with different options
+        # is not closed either.
+        client3.admin.command("ping")
+        disconnect("disconnect_reused_client_test_3")
+        # 3rd client is now closed:
+        if PYMONGO_VERSION >= (4,):
+            with pytest.raises(InvalidOperation):
+                client3.admin.command("ping")
+
     def test_disconnect_all(self):
         connections = mongoengine.connection._connections
         dbs = mongoengine.connection._dbs