enhance: add compaction lifecycle observability APIs

pymilvus/client/async_grpc_handler.py

@@ -800,10 +800,15 @@ async def get_persistent_segment_infos(
         collection_name: str,
         timeout: Optional[float] = None,
         context: Optional[CallContext] = None,
+        states: Optional[List[Union[int, str]]] = None,
         **kwargs,
     ) -> List[milvus_types.PersistentSegmentInfo]:
         check_pass_param(collection_name=collection_name, timeout=timeout)
-        req = Prepare.get_persistent_segment_info_request(collection_name)
+        req = Prepare.get_persistent_segment_info_request(
+            collection_name,
+            states=states,
+            db_name=context.get_db_name() if context else "",
+        )
         response = await self._async_stub.GetPersistentSegmentInfo(
             req, timeout=timeout, metadata=_api_level_md(context)
         )

pymilvus/client/grpc_handler.py

@@ -100,6 +100,31 @@
 logger = logging.getLogger(__name__)
 
 
+def _parse_compaction_plans(
+    response: milvus_types.GetCompactionPlansResponse,
+    compaction_id: int = 0,
+    collection_name: str = "",
+) -> CompactionPlans:
+    plans = CompactionPlans(compaction_id, response.state, collection_name=collection_name)
+    plans.plans = [
+        Plan(
+            list(merge.sources),
+            merge.target,
+            plan_id=merge.plan_id,
+            trigger_id=merge.trigger_id,
+            collection_id=merge.collection_id,
+            partition_id=merge.partition_id,
+            channel=merge.channel,
+            compaction_type=merge.type,
+            state=merge.state,
+            failure_reason=merge.failure_reason,
+            targets=list(merge.targets) or None,
+        )
+        for merge in response.mergeInfos
+    ]
+    return plans
+
+
 class ReconnectHandler:
     def __init__(self, conns: object, connection_name: str, kwargs: object) -> None:
         self.connection_name = connection_name
@@ -1550,7 +1575,10 @@ def get_query_segment_info(
         context: Optional[CallContext] = None,
         **kwargs,
     ) -> List[milvus_types.QuerySegmentInfo]:
-        req = Prepare.get_query_segment_info_request(collection_name)
+        req = Prepare.get_query_segment_info_request(
+            collection_name,
+            db_name=context.get_db_name() if context else "",
+        )
         response = self._stub.GetQuerySegmentInfo(
             req, timeout=timeout, metadata=_api_level_md(context)
         )
@@ -2224,9 +2252,14 @@ def get_persistent_segment_infos(
         collection_name: str,
         timeout: Optional[float] = None,
         context: Optional[CallContext] = None,
+        states: Optional[Iterable[Union[int, str]]] = None,
         **kwargs,
     ) -> List[milvus_types.PersistentSegmentInfo]:
-        req = Prepare.get_persistent_segment_info_request(collection_name)
+        req = Prepare.get_persistent_segment_info_request(
+            collection_name,
+            states=states,
+            db_name=context.get_db_name() if context else "",
+        )
         response = self._stub.GetPersistentSegmentInfo(
             req, timeout=timeout, metadata=_api_level_md(context)
         )
@@ -2532,11 +2565,25 @@ def get_compaction_plans(
         )
         check_status(response.status)
 
-        cp = CompactionPlans(compaction_id, response.state)
-
-        cp.plans = [Plan(m.sources, m.target) for m in response.mergeInfos]
+        return _parse_compaction_plans(response, compaction_id=compaction_id)
 
-        return cp
+    @retry_on_rpc_failure()
+    def get_compaction_tasks(
+        self,
+        collection_name: str,
+        timeout: Optional[float] = None,
+        context: Optional[CallContext] = None,
+        **kwargs,
+    ) -> CompactionPlans:
+        req = Prepare.get_compaction_tasks(
+            collection_name,
+            db_name=context.get_db_name() if context else "",
+        )
+        response = self._stub.GetCompactionStateWithPlans(
+            req, timeout=timeout, metadata=_api_level_md(context)
+        )
+        check_status(response.status)
+        return _parse_compaction_plans(response, collection_name=collection_name)
 
     @retry_on_rpc_failure()
     def get_replicas(

pymilvus/client/prepare.py

@@ -2292,8 +2292,17 @@ def get_collection_stats_request(cls, collection_name: str):
         return milvus_types.GetCollectionStatisticsRequest(collection_name=collection_name)
 
     @classmethod
-    def get_persistent_segment_info_request(cls, collection_name: str):
-        return milvus_types.GetPersistentSegmentInfoRequest(collectionName=collection_name)
+    def get_persistent_segment_info_request(
+        cls,
+        collection_name: str,
+        states: Optional[Iterable[Union[int, str]]] = None,
+        db_name: str = "",
+    ):
+        return milvus_types.GetPersistentSegmentInfoRequest(
+            dbName=db_name,
+            collectionName=collection_name,
+            states=states or [],
+        )
 
     @classmethod
     def get_flush_state_request(cls, segment_ids: List[int], collection_name: str, flush_ts: int):
@@ -2302,8 +2311,11 @@ def get_flush_state_request(cls, segment_ids: List[int], collection_name: str, f
         )
 
     @classmethod
-    def get_query_segment_info_request(cls, collection_name: str):
-        return milvus_types.GetQuerySegmentInfoRequest(collectionName=collection_name)
+    def get_query_segment_info_request(cls, collection_name: str, db_name: str = ""):
+        return milvus_types.GetQuerySegmentInfoRequest(
+            dbName=db_name,
+            collectionName=collection_name,
+        )
 
     @classmethod
     def flush_param(cls, collection_names: List[str]):
@@ -2498,6 +2510,16 @@ def get_compaction_state_with_plans(cls, compaction_id: int):
         request.compactionID = compaction_id
         return request
 
+    @classmethod
+    def get_compaction_tasks(cls, collection_name: str, db_name: str = ""):
+        if not isinstance(collection_name, str) or not collection_name:
+            raise ParamError(message=f"collection_name value {collection_name} is illegal")
+
+        return milvus_types.GetCompactionPlansRequest(
+            db_name=db_name,
+            collection_name=collection_name,
+        )
+
     @classmethod
     def get_replicas(cls, collection_id: int):
         if collection_id is None or not isinstance(collection_id, int):

pymilvus/client/types.py

@@ -1,6 +1,6 @@
 import logging
 import time
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from enum import IntEnum
 from typing import Any, ClassVar, Dict, List, Optional, TypeVar, Union
 
@@ -308,28 +308,58 @@ def __repr__(self) -> str:
 
 
 class Plan:
-    def __init__(self, sources: list, target: int) -> None:
+    def __init__(
+        self,
+        sources: list,
+        target: int,
+        *,
+        plan_id: int = 0,
+        trigger_id: int = 0,
+        collection_id: int = 0,
+        partition_id: int = 0,
+        channel: str = "",
+        compaction_type: str = "",
+        state: str = "",
+        failure_reason: str = "",
+        targets: Optional[List[int]] = None,
+    ) -> None:
         self.sources = sources
         self.target = target
+        self.plan_id = plan_id
+        self.task_id = plan_id
+        self.trigger_id = trigger_id
+        self.collection_id = collection_id
+        self.partition_id = partition_id
+        self.channel = channel
+        self.compaction_type = compaction_type
+        self.state = state
+        self.failure_reason = failure_reason
+        self.targets = list(targets) if targets is not None else ([target] if target else [])
 
     def __repr__(self) -> str:
         return f"""
 Plan:
+ - plan id: {self.plan_id}
+ - trigger id: {self.trigger_id}
+ - type: {self.compaction_type}
+ - state: {self.state}
  - sources: {self.sources}
- - target: {self.target}
+ - targets: {self.targets}
 """
 
 
 class CompactionPlans:
-    def __init__(self, compaction_id: int, state: int) -> None:
+    def __init__(self, compaction_id: int, state: int, collection_name: str = "") -> None:
         self.compaction_id = compaction_id
+        self.collection_name = collection_name
         self.state = State.new(state)
         self.plans = []
 
     def __repr__(self) -> str:
         return f"""
 Compaction Plans:
  - compaction id: {self.compaction_id}
+ - collection name: {self.collection_name}
  - state: {self.state.name}
  - plans: {self.plans}
  """
@@ -1390,6 +1420,9 @@ class SegmentInfo:
     state: common_pb2.SegmentState
     level: common_pb2.SegmentLevel
     storage_version: int
+    partition_id: int = field(default=0, init=False)
+    insert_channel: str = field(default="", init=False)
+    compaction_from: List[int] = field(default_factory=list, init=False)
 
     @property
     def state_name(self) -> str:
@@ -1408,13 +1441,16 @@ def __repr__(self) -> str:
             f"is_sorted={self.is_sorted}, "
             f"state='{self.state_name}', "
             f"level='{self.level_name}', "
-            f"storage_version={self.storage_version})"
+            f"storage_version={self.storage_version}, "
+            f"partition_id={self.partition_id}, "
+            f"insert_channel='{self.insert_channel}', "
+            f"compaction_from={self.compaction_from})"
         )
 
 
 @dataclass
 class LoadedSegmentInfo(SegmentInfo):
-    partition_id: int
+    partition_id: int = field()
     index_name: str
     index_id: int
     node_ids: List[int]

pymilvus/grpc_gen/common_pb2.pyi

@@ -142,6 +142,7 @@ class MsgType(int, metaclass=_enum_type_wrapper.EnumTypeWrapper):
     AlterCollectionFunction: _ClassVar[MsgType]
     DropCollectionFunction: _ClassVar[MsgType]
     TruncateCollection: _ClassVar[MsgType]
+    SplitShard: _ClassVar[MsgType]
     CreatePartition: _ClassVar[MsgType]
     DropPartition: _ClassVar[MsgType]
     HasPartition: _ClassVar[MsgType]
@@ -260,6 +261,15 @@ class MsgType(int, metaclass=_enum_type_wrapper.EnumTypeWrapper):
     RefreshExternalCollection: _ClassVar[MsgType]
     GetRefreshExternalCollectionProgress: _ClassVar[MsgType]
     ListRefreshExternalCollectionJobs: _ClassVar[MsgType]
+    CreateRowPolicy: _ClassVar[MsgType]
+    DropRowPolicy: _ClassVar[MsgType]
+    ListRowPolicies: _ClassVar[MsgType]
+    UpdateRowPolicy: _ClassVar[MsgType]
+    SetRLSPrincipalTags: _ClassVar[MsgType]
+    GetRLSPrincipalTags: _ClassVar[MsgType]
+    ListRLSPrincipals: _ClassVar[MsgType]
+    DeleteRLSPrincipalTags: _ClassVar[MsgType]
+    GetExportSnapshotState: _ClassVar[MsgType]
 
 class DslType(int, metaclass=_enum_type_wrapper.EnumTypeWrapper):
     __slots__ = ()
@@ -386,6 +396,10 @@ class ObjectPrivilege(int, metaclass=_enum_type_wrapper.EnumTypeWrapper):
     PrivilegeUnpinSnapshotData: _ClassVar[ObjectPrivilege]
     PrivilegeRestoreExternalSnapshot: _ClassVar[ObjectPrivilege]
     PrivilegeExportSnapshot: _ClassVar[ObjectPrivilege]
+    PrivilegeSkipRLS: _ClassVar[ObjectPrivilege]
+    PrivilegeViewRLS: _ClassVar[ObjectPrivilege]
+    PrivilegeManageRLS: _ClassVar[ObjectPrivilege]
+    PrivilegeImportBinlog: _ClassVar[ObjectPrivilege]
 
 class StateCode(int, metaclass=_enum_type_wrapper.EnumTypeWrapper):
     __slots__ = ()
@@ -538,6 +552,7 @@ AddCollectionFunction: MsgType
 AlterCollectionFunction: MsgType
 DropCollectionFunction: MsgType
 TruncateCollection: MsgType
+SplitShard: MsgType
 CreatePartition: MsgType
 DropPartition: MsgType
 HasPartition: MsgType
@@ -656,6 +671,15 @@ AlterCollectionSchema: MsgType
 RefreshExternalCollection: MsgType
 GetRefreshExternalCollectionProgress: MsgType
 ListRefreshExternalCollectionJobs: MsgType
+CreateRowPolicy: MsgType
+DropRowPolicy: MsgType
+ListRowPolicies: MsgType
+UpdateRowPolicy: MsgType
+SetRLSPrincipalTags: MsgType
+GetRLSPrincipalTags: MsgType
+ListRLSPrincipals: MsgType
+DeleteRLSPrincipalTags: MsgType
+GetExportSnapshotState: MsgType
 Dsl: DslType
 BoolExprV1: DslType
 UndefiedState: CompactionState
@@ -764,6 +788,10 @@ PrivilegePinSnapshotData: ObjectPrivilege
 PrivilegeUnpinSnapshotData: ObjectPrivilege
 PrivilegeRestoreExternalSnapshot: ObjectPrivilege
 PrivilegeExportSnapshot: ObjectPrivilege
+PrivilegeSkipRLS: ObjectPrivilege
+PrivilegeViewRLS: ObjectPrivilege
+PrivilegeManageRLS: ObjectPrivilege
+PrivilegeImportBinlog: ObjectPrivilege
 Initializing: StateCode
 Healthy: StateCode
 Abnormal: StateCode
@@ -1197,3 +1225,11 @@ class SearchAggregationSpec(_message.Message):
     sub_aggregation: SearchAggregationSpec
     search_size: int
     def __init__(self, fields: _Optional[_Iterable[str]] = ..., size: _Optional[int] = ..., metrics: _Optional[_Mapping[str, MetricAggSpec]] = ..., order: _Optional[_Iterable[_Union[OrderSpec, _Mapping]]] = ..., top_hits: _Optional[_Union[TopHitsSpec, _Mapping]] = ..., sub_aggregation: _Optional[_Union[SearchAggregationSpec, _Mapping]] = ..., search_size: _Optional[int] = ...) -> None: ...
+
+class IDRange(_message.Message):
+    __slots__ = ("begin", "end")
+    BEGIN_FIELD_NUMBER: _ClassVar[int]
+    END_FIELD_NUMBER: _ClassVar[int]
+    begin: int
+    end: int
+    def __init__(self, begin: _Optional[int] = ..., end: _Optional[int] = ...) -> None: ...