diff --git a/pymilvus/client/async_grpc_handler.py b/pymilvus/client/async_grpc_handler.py index e9ff31dd8..26581b86c 100644 --- a/pymilvus/client/async_grpc_handler.py +++ b/pymilvus/client/async_grpc_handler.py @@ -4,7 +4,7 @@ import socket import time from pathlib import Path -from typing import Any, Dict, List, Optional, Tuple, Union +from typing import Any, Dict, List, Optional, Sequence, Tuple, Union from urllib import parse import grpc @@ -800,10 +800,15 @@ async def get_persistent_segment_infos( collection_name: str, timeout: Optional[float] = None, context: Optional[CallContext] = None, + states: Optional[Sequence[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) ) diff --git a/pymilvus/client/grpc_handler.py b/pymilvus/client/grpc_handler.py index a28910b33..017441dbf 100644 --- a/pymilvus/client/grpc_handler.py +++ b/pymilvus/client/grpc_handler.py @@ -4,7 +4,7 @@ import threading import time from pathlib import Path -from typing import Any, Callable, Dict, Iterable, List, Mapping, Optional, Tuple, Union +from typing import Any, Callable, Dict, Iterable, List, Mapping, Optional, Sequence, Tuple, Union from urllib import parse import grpc @@ -99,6 +99,60 @@ logger = logging.getLogger(__name__) +# Keep the Plan API's names stable across the public protobuf enum migration. +_COMPACTION_TYPE_NAMES = { + common_pb2.CompactionTypeUndefined: "UndefinedCompaction", + common_pb2.CompactionTypeMerge: "MergeCompaction", + common_pb2.CompactionTypeMix: "MixCompaction", + common_pb2.CompactionTypeSingle: "SingleCompaction", + common_pb2.CompactionTypeMinor: "MinorCompaction", + common_pb2.CompactionTypeMajor: "MajorCompaction", + common_pb2.CompactionTypeLevel0Delete: "Level0DeleteCompaction", + common_pb2.CompactionTypeClustering: "ClusteringCompaction", + common_pb2.CompactionTypeSort: "SortCompaction", + common_pb2.CompactionTypePartitionKeySort: "PartitionKeySortCompaction", + common_pb2.CompactionTypeClusteringPartitionKeySort: "ClusteringPartitionKeySortCompaction", + common_pb2.CompactionTypeBumpSchemaVersion: "BumpSchemaVersionCompaction", +} +_COMPACTION_TASK_STATE_NAMES = { + common_pb2.CompactionTaskStateUnknown: "unknown", + common_pb2.CompactionTaskStateExecuting: "executing", + common_pb2.CompactionTaskStatePipelining: "pipelining", + common_pb2.CompactionTaskStateCompleted: "completed", + common_pb2.CompactionTaskStateFailed: "failed", + common_pb2.CompactionTaskStateTimeout: "timeout", + common_pb2.CompactionTaskStateAnalyzing: "analyzing", + common_pb2.CompactionTaskStateIndexing: "indexing", + common_pb2.CompactionTaskStateCleaned: "cleaned", + common_pb2.CompactionTaskStateMetaSaved: "meta_saved", + common_pb2.CompactionTaskStateStatistic: "statistic", +} + + +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=_COMPACTION_TYPE_NAMES.get(merge.type, f"unknown({merge.type})"), + state=_COMPACTION_TASK_STATE_NAMES.get(merge.state, f"unknown({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: @@ -1550,7 +1604,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 +2281,14 @@ def get_persistent_segment_infos( collection_name: str, timeout: Optional[float] = None, context: Optional[CallContext] = None, + states: Optional[Sequence[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 +2594,25 @@ def get_compaction_plans( ) check_status(response.status) - cp = CompactionPlans(compaction_id, response.state) + return _parse_compaction_plans(response, compaction_id=compaction_id) - cp.plans = [Plan(m.sources, m.target) for m in response.mergeInfos] - - 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( diff --git a/pymilvus/client/prepare.py b/pymilvus/client/prepare.py index ace1f4411..dba925257 100644 --- a/pymilvus/client/prepare.py +++ b/pymilvus/client/prepare.py @@ -3,7 +3,7 @@ import json import re import warnings -from typing import Any, Dict, Iterable, List, Mapping, Optional, Union +from typing import Any, Dict, Iterable, List, Mapping, Optional, Sequence, Union import numpy as np import orjson @@ -2292,8 +2292,41 @@ 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[Sequence[Union[int, str]]] = None, + db_name: str = "", + ): + if states is None: + normalized_states = [] + else: + if isinstance(states, (str, bytes)) or not isinstance(states, Sequence): + message = "states must be a sequence of SegmentState names or values" + raise ParamError(message=message) + if not states: + message = "states must not be empty; use None for the default state filter" + raise ParamError(message=message) + + valid_names = set(common_types.SegmentState.keys()) + valid_values = set(common_types.SegmentState.values()) + normalized_states = [] + for state in states: + if isinstance(state, str) and state in valid_names: + normalized_states.append(common_types.SegmentState.Value(state)) + elif ( + isinstance(state, int) and not isinstance(state, bool) and state in valid_values + ): + normalized_states.append(state) + else: + message = f"invalid SegmentState value: {state!r}" + raise ParamError(message=message) + + return milvus_types.GetPersistentSegmentInfoRequest( + dbName=db_name, + collectionName=collection_name, + states=normalized_states, + ) @classmethod def get_flush_state_request(cls, segment_ids: List[int], collection_name: str, flush_ts: int): @@ -2302,8 +2335,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 +2534,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): diff --git a/pymilvus/client/types.py b/pymilvus/client/types.py index bd6e06618..6abbc1890 100644 --- a/pymilvus/client/types.py +++ b/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,21 +308,50 @@ 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 > 0 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 = [] @@ -330,6 +359,7 @@ 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 = 0 + insert_channel: str = "" + compaction_from: List[int] = field(default_factory=list) @property def state_name(self) -> str: @@ -1408,13 +1441,19 @@ 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() + # Keep inherited defaults out of the legacy constructor's required positional fields. + insert_channel: str = field(default="", init=False) + compaction_from: List[int] = field(default_factory=list, init=False) index_name: str index_id: int node_ids: List[int] diff --git a/pymilvus/milvus_client/async_milvus_client.py b/pymilvus/milvus_client/async_milvus_client.py index b748d0c62..62a9aefb2 100644 --- a/pymilvus/milvus_client/async_milvus_client.py +++ b/pymilvus/milvus_client/async_milvus_client.py @@ -2,7 +2,7 @@ import copy import time import types -from typing import Dict, List, Optional, Type, Union +from typing import Dict, List, Optional, Sequence, Type, Union from pymilvus.client import type_info from pymilvus.client.abstract import AnnSearchRequest, BaseRanker @@ -1879,6 +1879,7 @@ async def list_persistent_segments( self, collection_name: str, timeout: Optional[float] = None, + states: Optional[Sequence[Union[int, str]]] = None, **kwargs, ) -> List[SegmentInfo]: """List persistent segments for a collection. @@ -1895,6 +1896,7 @@ async def list_persistent_segments( conn = await self._get_connection() infos = await conn.get_persistent_segment_infos( collection_name, + states=states, timeout=timeout, context=self._generate_call_context(**kwargs), **kwargs, @@ -1909,6 +1911,9 @@ async def list_persistent_segments( state=info.state, level=info.level, storage_version=info.storage_version, + partition_id=info.partitionID, + insert_channel=info.insert_channel, + compaction_from=list(info.compaction_from), ) for info in infos ] diff --git a/pymilvus/milvus_client/milvus_client.py b/pymilvus/milvus_client/milvus_client.py index 0a3d43baa..ed3fad7dd 100644 --- a/pymilvus/milvus_client/milvus_client.py +++ b/pymilvus/milvus_client/milvus_client.py @@ -1,7 +1,7 @@ import copy import logging import time -from typing import Callable, Dict, List, Optional, Union +from typing import Callable, Dict, List, Optional, Sequence, Union from pymilvus.client import type_info from pymilvus.client.abstract import AnnSearchRequest, BaseRanker @@ -2797,12 +2797,15 @@ def list_persistent_segments( self, collection_name: str, timeout: Optional[float] = None, + states: Optional[Sequence[Union[int, str]]] = None, **kwargs, ) -> List[SegmentInfo]: """List persistent segments for a collection. Args: collection_name (str): The name of the collection. + states (Optional[Sequence[Union[int, str]]]): Segment states to include; when omitted, + the server preserves the legacy persistent-state filter. timeout (Optional[float]): An optional duration of time in seconds to allow for the RPC. **kwargs: Additional arguments. @@ -2811,24 +2814,49 @@ def list_persistent_segments( """ infos = self._get_connection().get_persistent_segment_infos( collection_name, + states=states, timeout=timeout, context=self._generate_call_context(**kwargs), **kwargs, ) return [ SegmentInfo( - info.segmentID, - info.collectionID, - collection_name, - info.num_rows, - info.is_sorted, - info.state, - info.level, - info.storage_version, + 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, + partition_id=info.partitionID, + insert_channel=info.insert_channel, + compaction_from=list(info.compaction_from), ) for info in infos ] + def list_segments( + self, + collection_name: str, + states: Optional[Sequence[Union[int, str]]] = None, + timeout: Optional[float] = None, + **kwargs, + ) -> List[SegmentInfo]: + """List collection segments still retained in the requested lifecycle states. + + Dropped segment metadata is subject to server-side garbage collection, so callers that + need lineage history must poll and cache it during the retention window. + """ + if states is None: + states = ["Growing", "Sealed", "Flushing", "Flushed", "Importing", "Dropped"] + return self.list_persistent_segments( + collection_name, + states=states, + timeout=timeout, + **kwargs, + ) + def get_compaction_plans( self, job_id: int, @@ -2849,6 +2877,24 @@ def get_compaction_plans( job_id, timeout=timeout, context=self._generate_call_context(**kwargs), **kwargs ) + def list_compaction_tasks( + self, + collection_name: str, + timeout: Optional[float] = None, + **kwargs, + ) -> CompactionPlans: + """List all compaction tasks still retained for a collection. + + Terminal tasks are subject to server-side garbage collection and are not an audit log. + """ + validate_param("collection_name", collection_name, str) + return self._get_connection().get_compaction_tasks( + collection_name, + timeout=timeout, + context=self._generate_call_context(**kwargs), + **kwargs, + ) + def _is_collection_loaded(self, collection_name: str, timeout: Optional[float] = None) -> bool: state_dict = self.get_load_state(collection_name, timeout=timeout) return state_dict.get("state") == LoadState.Loaded diff --git a/tests/unit/grpc_handler/test_utility.py b/tests/unit/grpc_handler/test_utility.py index 56c551c84..d7f37929f 100644 --- a/tests/unit/grpc_handler/test_utility.py +++ b/tests/unit/grpc_handler/test_utility.py @@ -6,7 +6,8 @@ import pytest from pymilvus import AnnSearchRequest, RRFRanker from pymilvus.client.cache import GlobalCache -from pymilvus.exceptions import AmbiguousIndexName, MilvusException +from pymilvus.client.call_context import CallContext +from pymilvus.exceptions import AmbiguousIndexName, MilvusException, ParamError from pymilvus.grpc_gen import common_pb2 from pymilvus.grpc_gen import milvus_pb2 as milvus_types @@ -59,6 +60,120 @@ def test_get_compaction_plans(self, handler): result = handler.get_compaction_plans(123) assert result is not None + def test_get_compaction_tasks(self, handler): + handler._stub.GetCompactionStateWithPlans.return_value = ( + milvus_types.GetCompactionPlansResponse( + status=common_pb2.Status(error_code=common_pb2.Success), + state=common_pb2.Completed, + mergeInfos=[ + milvus_types.CompactionMergeInfo( + sources=[1, 2], + target=3, + plan_id=10, + trigger_id=20, + collection_id=30, + partition_id=40, + channel="ch", + type=common_pb2.CompactionTypeMix, + state=common_pb2.CompactionTaskStateCleaned, + failure_reason="DataNode reported compaction failure", + targets=[3, 4], + ), + milvus_types.CompactionMergeInfo(sources=[5, 6], target=-1), + milvus_types.CompactionMergeInfo(sources=[7, 8], target=9), + ], + ) + ) + result = handler.get_compaction_tasks("coll", context=CallContext(db_name="test_db")) + request = handler._stub.GetCompactionStateWithPlans.call_args.args[0] + assert request.db_name == "test_db" + assert request.collection_name == "coll" + assert result.collection_name == "coll" + assert result.plans[0].task_id == 10 + assert result.plans[0].targets == [3, 4] + assert result.plans[0].state == "cleaned" + assert result.plans[0].compaction_type == "MixCompaction" + assert result.plans[0].failure_reason == "DataNode reported compaction failure" + assert result.plans[1].targets == [] + assert result.plans[1].state == "unknown" + assert result.plans[1].compaction_type == "UndefinedCompaction" + assert result.plans[2].targets == [9] + + @pytest.mark.parametrize("by_collection", [False, True]) + @pytest.mark.parametrize( + "task_type,type_name", + [ + (0, "UndefinedCompaction"), + (2, "MergeCompaction"), + (3, "MixCompaction"), + (4, "SingleCompaction"), + (5, "MinorCompaction"), + (6, "MajorCompaction"), + (7, "Level0DeleteCompaction"), + (8, "ClusteringCompaction"), + (9, "SortCompaction"), + (10, "PartitionKeySortCompaction"), + (11, "ClusteringPartitionKeySortCompaction"), + (12, "BumpSchemaVersionCompaction"), + (99, "unknown(99)"), + ], + ) + @pytest.mark.parametrize( + "task_state,state_name", + [ + *enumerate( + [ + "unknown", + "executing", + "pipelining", + "completed", + "failed", + "timeout", + "analyzing", + "indexing", + "cleaned", + "meta_saved", + "statistic", + ] + ), + (99, "unknown(99)"), + ], + ) + def test_compaction_enum_wire_roundtrip( + self, handler, by_collection, task_type, type_name, task_state, state_name + ): + response = milvus_types.GetCompactionPlansResponse( + status=common_pb2.Status(error_code=common_pb2.Success), + state=common_pb2.Completed, + mergeInfos=[ + milvus_types.CompactionMergeInfo( + plan_id=10, + sources=[1, 2], + target=3, + targets=[3, 4], + type=task_type, + state=task_state, + failure_reason="retained failure reason", + ) + ], + ) + # Exercise the real enum wire types, not mocked string attributes. + handler._stub.GetCompactionStateWithPlans.return_value = ( + milvus_types.GetCompactionPlansResponse.FromString(response.SerializeToString()) + ) + result = ( + handler.get_compaction_tasks("coll") + if by_collection + else handler.get_compaction_plans(123) + ) + assert result.state.name == "Completed" + assert result.plans[0].compaction_type == type_name + assert result.plans[0].state == state_name + assert result.plans[0].failure_reason == "retained failure reason" + assert result.plans[0].sources == [1, 2] + assert result.plans[0].targets == [3, 4] + assert result.plans[0].target == 3 + def test_get_server_version(self, handler): handler._stub.GetVersion.return_value = make_response(version="v2.4.0") assert handler.get_server_version() == "v2.4.0" @@ -182,14 +297,33 @@ class TestGrpcHandlerSegmentOps: def test_get_query_segment_info(self, handler): mock_seg = MagicMock(segmentID=1, collectionID=100) handler._stub.GetQuerySegmentInfo.return_value = make_response(infos=[mock_seg]) - handler.get_query_segment_info("coll") + handler.get_query_segment_info("coll", context=CallContext(db_name="test_db")) handler._stub.GetQuerySegmentInfo.assert_called_once() + request = handler._stub.GetQuerySegmentInfo.call_args.args[0] + assert request.dbName == "test_db" def test_get_persistent_segment_infos(self, handler): mock_seg = MagicMock(segmentID=1, num_rows=1000) handler._stub.GetPersistentSegmentInfo.return_value = make_response(infos=[mock_seg]) - handler.get_persistent_segment_infos("coll") + handler.get_persistent_segment_infos( + "coll", + states=["Growing", "Dropped"], + context=CallContext(db_name="test_db"), + ) handler._stub.GetPersistentSegmentInfo.assert_called_once() + request = handler._stub.GetPersistentSegmentInfo.call_args.args[0] + assert request.dbName == "test_db" + assert list(request.states) == [ + common_pb2.SegmentState.Growing, + common_pb2.SegmentState.Dropped, + ] + + @pytest.mark.parametrize("states", ["Flushed", iter(["Dropped"]), [], [99]]) + def test_get_persistent_segment_infos_rejects_invalid_states_before_rpc(self, handler, states): + with pytest.raises(ParamError): + handler.get_persistent_segment_infos("coll", states=states) + + handler._stub.GetPersistentSegmentInfo.assert_not_called() class TestGrpcHandlerImportExport: diff --git a/tests/unit/prepare/test_coverage_gaps.py b/tests/unit/prepare/test_coverage_gaps.py index 6d597432a..a926e8bd6 100644 --- a/tests/unit/prepare/test_coverage_gaps.py +++ b/tests/unit/prepare/test_coverage_gaps.py @@ -3,7 +3,8 @@ import numpy as np import pytest from pymilvus.client.prepare import Prepare -from pymilvus.exceptions import ParamError +from pymilvus.exceptions import ErrorCode, ParamError +from pymilvus.grpc_gen import common_pb2 class TestCreateCollectionNumPartitions: @@ -259,13 +260,66 @@ class TestSegmentRequests: def test_get_persistent_segment_info(self): """Test get persistent segment info request.""" - req = Prepare.get_persistent_segment_info_request("test_coll") + req = Prepare.get_persistent_segment_info_request( + "test_coll", states=["Growing", "Dropped"], db_name="test_db" + ) assert req.collectionName == "test_coll" + assert req.dbName == "test_db" + assert list(req.states) == [ + common_pb2.SegmentState.Growing, + common_pb2.SegmentState.Dropped, + ] + + def test_get_persistent_segment_info_reuses_state_sequence(self): + states = ("Growing", common_pb2.SegmentState.Dropped) + + first = Prepare.get_persistent_segment_info_request("test_coll", states=states) + second = Prepare.get_persistent_segment_info_request("test_coll", states=states) + + expected = [common_pb2.SegmentState.Growing, common_pb2.SegmentState.Dropped] + assert list(first.states) == expected + assert list(second.states) == expected + + @pytest.mark.parametrize( + "states", + [ + "Flushed", + iter(["Dropped"]), + {"Dropped"}, + ], + ) + def test_get_persistent_segment_info_rejects_non_sequence_states(self, states): + with pytest.raises(ParamError, match="states must be a sequence") as exc_info: + Prepare.get_persistent_segment_info_request("test_coll", states=states) + assert exc_info.value.code == ErrorCode.UNEXPECTED_ERROR + assert exc_info.value.message == "states must be a sequence of SegmentState names or values" + + def test_get_persistent_segment_info_rejects_empty_states(self): + with pytest.raises(ParamError, match="states must not be empty") as exc_info: + Prepare.get_persistent_segment_info_request("test_coll", states=[]) + assert exc_info.value.code == ErrorCode.UNEXPECTED_ERROR + assert exc_info.value.message == ( + "states must not be empty; use None for the default state filter" + ) + + @pytest.mark.parametrize("state", ["NotAState", 99, True, None]) + def test_get_persistent_segment_info_rejects_invalid_state(self, state): + with pytest.raises(ParamError, match="invalid SegmentState value") as exc_info: + Prepare.get_persistent_segment_info_request("test_coll", states=[state]) + assert exc_info.value.code == ErrorCode.UNEXPECTED_ERROR + assert exc_info.value.message == f"invalid SegmentState value: {state!r}" + + def test_get_compaction_tasks(self): + req = Prepare.get_compaction_tasks("test_coll", db_name="test_db") + assert req.collection_name == "test_coll" + assert req.db_name == "test_db" + assert req.compactionID == 0 def test_get_query_segment_info(self): """Test get query segment info request.""" - req = Prepare.get_query_segment_info_request("test_coll") + req = Prepare.get_query_segment_info_request("test_coll", db_name="test_db") assert req.collectionName == "test_coll" + assert req.dbName == "test_db" class TestPartitionRequests: diff --git a/tests/unit/test_async_milvus_client.py b/tests/unit/test_async_milvus_client.py index 574d6c0e4..5fa682b64 100644 --- a/tests/unit/test_async_milvus_client.py +++ b/tests/unit/test_async_milvus_client.py @@ -163,11 +163,14 @@ async def test_list_persistent_segments(self): mock_segment_info = MagicMock() mock_segment_info.segmentID = 1001 mock_segment_info.collectionID = 2001 + mock_segment_info.partitionID = 3001 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_segment_info.insert_channel = "test-channel" + mock_segment_info.compaction_from = [10, 11] mock_handler.get_persistent_segment_infos = AsyncMock(return_value=[mock_segment_info]) @@ -188,15 +191,20 @@ async def test_list_persistent_segments(self): assert segment_info.segment_id == 1001 assert segment_info.collection_id == 2001 assert segment_info.collection_name == "test_collection" + assert segment_info.partition_id == 3001 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 + assert segment_info.insert_channel == "test-channel" + assert segment_info.compaction_from == [10, 11] + segment_info.compaction_from.append(12) + assert mock_segment_info.compaction_from == [10, 11] # Verify call arguments mock_handler.get_persistent_segment_infos.assert_called_once_with( - "test_collection", timeout=None, context=ANY + "test_collection", states=None, timeout=None, context=ANY ) @pytest.mark.asyncio diff --git a/tests/unit/test_client_types.py b/tests/unit/test_client_types.py index 40b4abe26..a0b6949fc 100644 --- a/tests/unit/test_client_types.py +++ b/tests/unit/test_client_types.py @@ -265,9 +265,18 @@ def test_compaction_state_repr(self): # TestPlan class TestPlan: def test_plan_init(self): - plan = Plan(sources=[1, 2, 3], target=100) + plan = Plan( + sources=[1, 2, 3], + target=100, + plan_id=10, + state="completed", + targets=[100, 101], + ) assert plan.sources == [1, 2, 3] assert plan.target == 100 + assert plan.task_id == 10 + assert plan.state == "completed" + assert plan.targets == [100, 101] def test_plan_repr(self): r = repr(Plan(sources=[10, 20], target=200)) diff --git a/tests/unit/test_milvus_client.py b/tests/unit/test_milvus_client.py index 3fa3a5916..31b5c6cb8 100644 --- a/tests/unit/test_milvus_client.py +++ b/tests/unit/test_milvus_client.py @@ -1915,9 +1915,58 @@ def test_list_loaded_segments(self, mc): assert client.list_loaded_segments("col") == [] def test_list_persistent_segments(self, mc): + client, handler = mc + segment = MagicMock( + segmentID=1, + collectionID=2, + partitionID=4, + num_rows=3, + is_sorted=True, + state=common_pb2.SegmentState.Flushed, + level=common_pb2.SegmentLevel.L1, + storage_version=3, + insert_channel="ch", + compaction_from=[10, 11], + ) + handler.get_persistent_segment_infos.return_value = [segment] + result = client.list_persistent_segments("col", states=["Flushed", "Dropped"]) + assert len(result) == 1 + assert result[0].segment_id == 1 + assert result[0].collection_id == 2 + assert result[0].collection_name == "col" + assert result[0].num_rows == 3 + assert result[0].is_sorted is True + assert result[0].state == common_pb2.SegmentState.Flushed + assert result[0].level == common_pb2.SegmentLevel.L1 + assert result[0].storage_version == 3 + assert result[0].partition_id == 4 + assert result[0].insert_channel == "ch" + assert result[0].compaction_from == [10, 11] + result[0].compaction_from.append(12) + assert segment.compaction_from == [10, 11] + handler.get_persistent_segment_infos.assert_called_once_with( + "col", states=["Flushed", "Dropped"], timeout=None, context=ANY + ) + + def test_list_segments_defaults_to_all_lifecycle_states(self, mc): client, handler = mc handler.get_persistent_segment_infos.return_value = [] - assert client.list_persistent_segments("col") == [] + assert client.list_segments("col") == [] + assert handler.get_persistent_segment_infos.call_args.kwargs["states"] == [ + "Growing", + "Sealed", + "Flushing", + "Flushed", + "Importing", + "Dropped", + ] + + def test_list_compaction_tasks(self, mc): + client, handler = mc + expected = MagicMock() + handler.get_compaction_tasks.return_value = expected + assert client.list_compaction_tasks("col") is expected + handler.get_compaction_tasks.assert_called_once_with("col", timeout=None, context=ANY) def test_using_database(self, mc): client, handler = mc diff --git a/tests/unit/test_types.py b/tests/unit/test_types.py index 75b7b7116..78544027c 100644 --- a/tests/unit/test_types.py +++ b/tests/unit/test_types.py @@ -94,6 +94,65 @@ def test_shard_dup_nodeIDs(self): assert s.shard_leader == 1 +class TestSegmentInfoConstruction: + @pytest.fixture + def segment_kwargs(self): + return { + "segment_id": 123, + "collection_id": 456, + "collection_name": "test_col", + "num_rows": 1000, + "is_sorted": True, + "state": 4, + "level": 2, + "storage_version": 3, + } + + @pytest.mark.parametrize("use_keywords", [False, True]) + @pytest.mark.parametrize( + "segment_type,extra_kwargs", + [ + (SegmentInfo, {}), + ( + LoadedSegmentInfo, + { + "partition_id": 100, + "index_name": "idx_vec", + "index_id": 200, + "node_ids": [1, 2], + "mem_size": 4096, + }, + ), + ], + ) + def test_legacy_constructor(self, segment_kwargs, segment_type, extra_kwargs, use_keywords): + kwargs = {**segment_kwargs, **extra_kwargs} + info = segment_type(**kwargs) if use_keywords else segment_type(*kwargs.values()) + + for name, expected in kwargs.items(): + assert getattr(info, name) == expected + assert info.partition_id == extra_kwargs.get("partition_id", 0) + assert info.insert_channel == "" + assert info.compaction_from == [] + + other = segment_type(**kwargs) + info.compaction_from.append(10) + assert other.compaction_from == [] + + @pytest.mark.parametrize("use_keywords", [False, True]) + def test_constructor_accepts_lifecycle_metadata(self, segment_kwargs, use_keywords): + kwargs = { + **segment_kwargs, + "partition_id": 100, + "insert_channel": "test-channel", + "compaction_from": [10, 11], + } + info = SegmentInfo(**kwargs) if use_keywords else SegmentInfo(*kwargs.values()) + + for name, expected in kwargs.items(): + assert getattr(info, name) == expected + + class TestSegmentInfoRepr: def test_segment_info_repr_shows_state_and_level_names(self): info = SegmentInfo(