[Backport 2.6] feat:support async client list_persistent_segments (#3206) (#3207)

Backport of #3206 to `2.6`.

Signed-off-by: silas.jiang <silas.jiang@zilliz.com>
Signed-off-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
Co-authored-by: jac <jacllovey@qq.com>
Co-authored-by: silas.jiang <silas.jiang@zilliz.com>

Author: 3 people

pymilvus/client/async_grpc_handler.py

@@ -628,6 +628,18 @@ async def _prepare_row_insert_request(
         )
 
     @retry_on_rpc_failure()
+    async def get_persistent_segment_infos(
+        self, collection_name: str, timeout: Optional[float] = None, **kwargs
+    ) -> List[milvus_types.PersistentSegmentInfo]:
+        await self.ensure_channel_ready()
+        check_pass_param(collection_name=collection_name, timeout=timeout)
+        req = Prepare.get_persistent_segment_info_request(collection_name)
+        response = await self._async_stub.GetPersistentSegmentInfo(
+            req, timeout=timeout, metadata=_api_level_md(**kwargs)
+        )
+        check_status(response.status)
+        return response.infos
+
     async def delete(
         self,
         collection_name: str,

pymilvus/milvus_client/async_milvus_client.py

@@ -8,6 +8,7 @@
     OmitZeroDict,
     ResourceGroupConfig,
     RoleInfo,
+    SegmentInfo,
     UserInfo,
 )
 from pymilvus.client.utils import convert_struct_fields_to_user_format, is_vector_type
@@ -1154,6 +1155,40 @@ async def get_flush_all_state(self, timeout: Optional[float] = None, **kwargs) -
         conn = self._get_connection()
         return await conn.get_flush_all_state(timeout=timeout, **kwargs)
 
+    async def list_persistent_segments(
+        self,
+        collection_name: str,
+        timeout: Optional[float] = None,
+        **kwargs,
+    ) -> List[SegmentInfo]:
+        """List persistent segments for a collection.
+
+        Args:
+            collection_name (str): The name of the collection.
+            timeout (Optional[float]): An optional duration of time in seconds to allow for the RPC.
+            **kwargs: Additional arguments.
+
+        Returns:
+            List[SegmentInfo]: A list of persistent segment information.
+        """
+        validate_param("collection_name", collection_name, str)
+        infos = await self._get_connection().get_persistent_segment_infos(
+            collection_name, timeout=timeout, **kwargs
+        )
+        return [
+            SegmentInfo(
+                segment_id=info.segmentID,
+                collection_id=info.collectionID,
+                collection_name=collection_name,
+                num_rows=info.num_rows,
+                is_sorted=info.is_sorted,
+                state=info.state,
+                level=info.level,
+                storage_version=info.storage_version,
+            )
+            for info in infos
+        ]
+
     async def compact(
         self,
         collection_name: str,

tests/test_async_milvus_client.py

@@ -86,6 +86,49 @@ async def test_describe_collection_with_struct_array_fields(self):
             assert len(result["fields"]) == 2  # original field + converted struct field
             mock_handler.describe_collection.assert_called_once()
 
+    @pytest.mark.asyncio
+    async def test_list_persistent_segments(self):
+        with patch('pymilvus.milvus_client.async_milvus_client.create_connection', return_value="test"), \
+             patch('pymilvus.orm.connections.Connections._fetch_handler') as mock_fetch:
+            # Mock connection and its get_persistent_segment_infos method
+            mock_conn = AsyncMock()
+            mock_fetch.return_value = mock_conn
+            
+            # Create a mock segment info
+            mock_segment_info = MagicMock()
+            mock_segment_info.segmentID = 1001
+            mock_segment_info.collectionID = 2001
+            mock_segment_info.num_rows = 1000
+            mock_segment_info.is_sorted = True
+            mock_segment_info.state = 3  # FLUSHED
+            mock_segment_info.level = 1
+            mock_segment_info.storage_version = 1
+            
+            mock_conn.get_persistent_segment_infos.return_value = [mock_segment_info]
+
+            # Initialize AsyncMilvusClient
+            client = AsyncMilvusClient(uri="http://localhost:19530")
+
+            # Call list_persistent_segments
+            result = await client.list_persistent_segments("test_collection")
+
+            # Verify the result
+            assert len(result) == 1
+            segment_info = result[0]
+            assert segment_info.segment_id == 1001
+            assert segment_info.collection_id == 2001
+            assert segment_info.collection_name == "test_collection"
+            assert segment_info.num_rows == 1000
+            assert segment_info.is_sorted is True
+            assert segment_info.state == 3
+            assert segment_info.level == 1
+            assert segment_info.storage_version == 1
+
+            # Verify call arguments
+            mock_conn.get_persistent_segment_infos.assert_called_once_with(
+                "test_collection", timeout=None
+            )
+
     @pytest.mark.asyncio
     async def test_describe_collection_without_struct_array_fields(self):
         """Test describe_collection works normally when no struct_array_fields"""