<ADD>: add disconnect_all_async

Author: jeonghanjoo

mongoengine/connection.py

@@ -39,6 +39,7 @@
     "disconnect",
     "disconnect_async",
     "disconnect_all",
+    "disconnect_all_async",
     "get_connection",
     "get_db",
     "get_async_db",
@@ -685,11 +686,8 @@ async def disconnect_async(alias=DEFAULT_CONNECTION_NAME):
     if connection:
         # Only close if this is the last reference to this connection
         if all(connection is not c for c in _connections.values()):
-            if is_async_connection(alias):
-                # For AsyncMongoClient, we need to call close() method
-                connection.close()
-            else:
-                connection.close()
+            # AsyncMongoClient.close() is a coroutine, must be awaited
+            await connection.close()
 
     # Clean up database references
     if alias in _dbs:
@@ -708,6 +706,23 @@ async def disconnect_async(alias=DEFAULT_CONNECTION_NAME):
         del _connection_types[alias]
 
 
+async def disconnect_all_async():
+    """Close all registered async database connections.
+
+    This is the async version of disconnect_all() that properly closes
+    AsyncMongoClient connections. It will only close async connections,
+    leaving sync connections untouched.
+    """
+    # Get list of all async connections
+    async_aliases = [
+        alias for alias in list(_connections.keys()) if is_async_connection(alias)
+    ]
+
+    # Disconnect each async connection
+    for alias in async_aliases:
+        await disconnect_async(alias)
+
+
 def get_async_db(alias=DEFAULT_CONNECTION_NAME, reconnect=False):
     """Get the async database for a given alias.
 

tests/test_async_connection.py

@@ -7,6 +7,7 @@
     connect,
     connect_async,
     disconnect,
+    disconnect_all_async,
     disconnect_async,
     get_async_db,
     is_async_connection,
@@ -178,3 +179,71 @@ async def test_reconnect_async_different_settings(self):
 
         # Clean up
         await disconnect_async("reconnect_test2")
+
+    @pytest.mark.asyncio
+    async def test_disconnect_all_async(self):
+        """Test disconnect_all_async only disconnects async connections."""
+        # Create mix of sync and async connections
+        connect(db="sync_db1", alias="sync1")
+        connect(db="sync_db2", alias="sync2")
+        await connect_async(db="async_db1", alias="async1")
+        await connect_async(db="async_db2", alias="async2")
+        await connect_async(db="async_db3", alias="async3")
+
+        # Verify connections exist
+        assert not is_async_connection("sync1")
+        assert not is_async_connection("sync2")
+        assert is_async_connection("async1")
+        assert is_async_connection("async2")
+        assert is_async_connection("async3")
+
+        from mongoengine.connection import _connections
+
+        assert len(_connections) == 5
+
+        # Disconnect all async connections
+        await disconnect_all_async()
+
+        # Verify only async connections were disconnected
+        assert "sync1" in _connections
+        assert "sync2" in _connections
+        assert "async1" not in _connections
+        assert "async2" not in _connections
+        assert "async3" not in _connections
+
+        # Verify sync connections still work
+        assert not is_async_connection("sync1")
+        assert not is_async_connection("sync2")
+
+        # Clean up remaining sync connections
+        disconnect("sync1")
+        disconnect("sync2")
+
+    @pytest.mark.asyncio
+    async def test_disconnect_all_async_empty(self):
+        """Test disconnect_all_async when no connections exist."""
+        # Should not raise any errors
+        await disconnect_all_async()
+
+    @pytest.mark.asyncio
+    async def test_disconnect_all_async_only_sync(self):
+        """Test disconnect_all_async when only sync connections exist."""
+        # Create only sync connections
+        connect(db="sync_db1", alias="sync1")
+        connect(db="sync_db2", alias="sync2")
+
+        from mongoengine.connection import _connections
+
+        assert len(_connections) == 2
+
+        # Disconnect all async (should do nothing)
+        await disconnect_all_async()
+
+        # Verify sync connections still exist
+        assert len(_connections) == 2
+        assert "sync1" in _connections
+        assert "sync2" in _connections
+
+        # Clean up
+        disconnect("sync1")
+        disconnect("sync2")