milvus/tests/python_client/base/client_v2_base.py

import sys
import time

from pymilvus import DataType, MilvusClient

sys.path.append("..")
from base.client_base import Base
from check.func_check import ResponseChecker
from common import common_func as cf
from common import common_type as ct
from utils.api_request import api_request
from utils.util_log import test_log as log
from utils.wrapper import trace

TIMEOUT = 120
INDEX_NAME = ""


class TestMilvusClientV2Base(Base):
    # milvus_client = None
    active_trace = False

    def init_async_milvus_client(self):
        uri = cf.param_info.param_uri or f"http://{cf.param_info.param_host}:{cf.param_info.param_port}"
        kwargs = {
            "uri": uri,
            "user": cf.param_info.param_user,
            "password": cf.param_info.param_password,
            "token": cf.param_info.param_token,
        }
        self.async_milvus_client_wrap.init_async_client(**kwargs)

    def _client(self, active_trace=False, **kwargs):
        """return MilvusClient instance if connected successfully, otherwise return None"""
        if self.skip_connection:
            return None
        if cf.param_info.param_uri:
            uri = cf.param_info.param_uri
        else:
            uri = "http://" + cf.param_info.param_host + ":" + str(cf.param_info.param_port)
        res, is_succ = self.init_milvus_client(
            uri=uri, token=cf.param_info.param_token, active_trace=active_trace, **kwargs
        )
        if is_succ:
            # self.milvus_client = res
            log.info(f"server version: {res.get_server_version(detail=True)}")
        return res

    def init_milvus_client(
        self,
        uri,
        user="",
        password="",
        db_name="",
        token="",
        timeout=None,
        check_task=None,
        check_items=None,
        active_trace=False,
        **kwargs,
    ):
        self.active_trace = active_trace
        func_name = sys._getframe().f_code.co_name
        res, is_succ = api_request([MilvusClient, uri, user, password, db_name, token, timeout], **kwargs)
        # self.milvus_client = res if is_succ else None
        check_result = ResponseChecker(
            res,
            func_name,
            check_task,
            check_items,
            is_succ,
            uri=uri,
            user=user,
            password=password,
            db_name=db_name,
            token=token,
            timeout=timeout,
            **kwargs,
        ).run()
        return res, check_result

    @trace()
    def close(self, client, check_task=None, check_items=None):
        func_name = sys._getframe().f_code.co_name
        res, is_succ = api_request([client.close])
        check_result = ResponseChecker(res, func_name, check_task, check_items, is_succ).run()
        return res, check_result

    @trace()
    def create_schema(self, client, timeout=None, check_task=None, check_items=None, **kwargs):
        timeout = TIMEOUT if timeout is None else timeout
        kwargs.update({"timeout": timeout})

        func_name = sys._getframe().f_code.co_name
        res, check = api_request([client.create_schema], **kwargs)
        check_result = ResponseChecker(res, func_name, check_task, check_items, check, **kwargs).run()
        return res, check_result

    @trace()
    def create_struct_field_schema(self, client, check_task=None, check_items=None, **kwargs):

        func_name = sys._getframe().f_code.co_name
        res, check = api_request([client.create_struct_field_schema], **kwargs)
        check_result = ResponseChecker(res, func_name, check_task, check_items, check, **kwargs).run()
        return res, check_result

    @trace()
    def add_field(self, schema, field_name, datatype, check_task=None, check_items=None, **kwargs):

        # Set default parameters for specific field types
        if datatype == DataType.VARCHAR and "max_length" not in kwargs:
            kwargs["max_length"] = ct.default_length
        elif datatype == DataType.ARRAY:
            if "element_type" not in kwargs:
                kwargs["element_type"] = DataType.INT64
            if "max_capacity" not in kwargs:
                kwargs["max_capacity"] = ct.default_max_capacity

        func_name = sys._getframe().f_code.co_name
        res, check = api_request([schema.add_field, field_name, datatype], **kwargs)
        check_result = ResponseChecker(res, func_name, check_task, check_items, check, **kwargs).run()
        return res, check_result

    @trace()
    def create_collection(
        self,
        client,
        collection_name,
        dimension=None,
        primary_field_name="id",
        id_type="int",
        vector_field_name="vector",
        metric_type="COSINE",
        auto_id=False,
        schema=None,
        index_params=None,
        timeout=None,
        force_teardown=True,
        check_task=None,
        check_items=None,
        **kwargs,
    ):
        timeout = TIMEOUT if timeout is None else timeout
        consistency_level = kwargs.get("consistency_level", "Strong")
        kwargs.update({"consistency_level": consistency_level})

        func_name = sys._getframe().f_code.co_name
        res, check = api_request(
            [
                client.create_collection,
                collection_name,
                dimension,
                primary_field_name,
                id_type,
                vector_field_name,
                metric_type,
                auto_id,
                timeout,
                schema,
                index_params,
            ],
            **kwargs,
        )
        # Register for cleanup BEFORE assertion check, so collection is cleaned up even if check fails
        # (e.g., when check_task=err_res but API actually succeeds due to server bug)
        if check and force_teardown:
            self.tear_down_collection_names.append(collection_name)
        check_result = ResponseChecker(
            res,
            func_name,
            check_task,
            check_items,
            check,
            collection_name=collection_name,
            dimension=dimension,
            **kwargs,
        ).run()
        return res, check_result

    def has_collection(self, client, collection_name, timeout=None, check_task=None, check_items=None, **kwargs):
        timeout = TIMEOUT if timeout is None else timeout
        kwargs.update({"timeout": timeout})

        func_name = sys._getframe().f_code.co_name
        res, check = api_request([client.has_collection, collection_name], **kwargs)
        check_result = ResponseChecker(
            res, func_name, check_task, check_items, check, collection_name=collection_name, **kwargs
        ).run()
        return res, check_result

    @trace()
    def insert(self, client, collection_name, data, timeout=None, check_task=None, check_items=None, **kwargs):
        timeout = TIMEOUT if timeout is None else timeout
        kwargs.update({"timeout": timeout})

        func_name = sys._getframe().f_code.co_name
        res, check = api_request([client.insert, collection_name, data], **kwargs)
        check_result = ResponseChecker(res, func_name, check_task, check_items, check, **kwargs).run()
        return res, check_result

    @trace()
    def upsert(self, client, collection_name, data, timeout=None, check_task=None, check_items=None, **kwargs):
        timeout = TIMEOUT if timeout is None else timeout
        kwargs.update({"timeout": timeout})

        func_name = sys._getframe().f_code.co_name
        res, check = api_request([client.upsert, collection_name, data], **kwargs)
        check_result = ResponseChecker(
            res, func_name, check_task, check_items, check, collection_name=collection_name, data=data, **kwargs
        ).run()
        return res, check_result

    @trace()
    def get_collection_stats(self, client, collection_name, timeout=None, check_task=None, check_items=None, **kwargs):
        timeout = TIMEOUT if timeout is None else timeout
        kwargs.update({"timeout": timeout})

        func_name = sys._getframe().f_code.co_name
        res, check = api_request([client.get_collection_stats, collection_name], **kwargs)
        check_result = ResponseChecker(
            res, func_name, check_task, check_items, check, collection_name=collection_name, **kwargs
        ).run()
        return res, check_result

    @trace()
    def search(
        self,
        client,
        collection_name,
        data=None,
        limit=10,
        filter=None,
        output_fields=None,
        search_params=None,
        timeout=None,
        check_task=None,
        check_items=None,
        **kwargs,
    ):
        timeout = TIMEOUT if timeout is None else timeout
        # kwargs.update({"timeout": timeout})

        func_name = sys._getframe().f_code.co_name
        res, check = api_request(
            [client.search, collection_name, data, filter, limit, output_fields, search_params], **kwargs
        )
        check_result = ResponseChecker(
            res,
            func_name,
            check_task,
            check_items,
            check,
            collection_name=collection_name,
            data=data,
            limit=limit,
            filter=filter,
            output_fields=output_fields,
            search_params=search_params,
            **kwargs,
        ).run()
        return res, check_result

    @trace()
    def search_iterator(
        self,
        client,
        collection_name,
        data,
        batch_size,
        limit=-1,
        filter=None,
        output_fields=None,
        search_params=None,
        timeout=None,
        check_task=None,
        check_items=None,
        **kwargs,
    ):
        timeout = TIMEOUT if timeout is None else timeout
        kwargs.update({"timeout": timeout})

        func_name = sys._getframe().f_code.co_name
        res, check = api_request(
            [client.search_iterator, collection_name, data, batch_size, filter, limit, output_fields, search_params],
            **kwargs,
        )
        if any(k in kwargs for k in ["use_rbac_mul_db", "use_mul_db"]):
            self.using_database(client, kwargs.get("another_db"))
        if kwargs.get("use_alias", False) is True:
            alias = collection_name
            self.alter_alias(client, kwargs.get("another_collection"), alias)
        check_result = ResponseChecker(
            res,
            func_name,
            check_task,
            check_items,
            check,
            collection_name=collection_name,
            data=data,
            batch_size=batch_size,
            filter=filter,
            limit=limit,
            output_fields=output_fields,
            search_params=search_params,
            **kwargs,
        ).run()
        return res, check_result

    @trace()
    def hybrid_search(
        self,
        client,
        collection_name,
        reqs,
        ranker,
        limit=10,
        output_fields=None,
        timeout=None,
        partition_names=None,
        check_task=None,
        check_items=None,
        **kwargs,
    ):
        timeout = TIMEOUT if timeout is None else timeout
        # kwargs.update({"timeout": timeout})
        func_name = sys._getframe().f_code.co_name
        res, check = api_request(
            [client.hybrid_search, collection_name, reqs, ranker, limit, output_fields, timeout, partition_names],
            **kwargs,
        )
        check_result = ResponseChecker(
            res,
            func_name,
            check_task,
            check_items,
            check,
            collection_name=collection_name,
            reqs=reqs,
            ranker=ranker,
            limit=limit,
            output_fields=output_fields,
            timeout=timeout,
            partition_names=partition_names,
            **kwargs,
        ).run()
        return res, check_result

    @trace()
    def query(self, client, collection_name, timeout=None, check_task=None, check_items=None, **kwargs):
        timeout = TIMEOUT if timeout is None else timeout
        kwargs.update({"timeout": timeout})

        func_name = sys._getframe().f_code.co_name
        res, check = api_request([client.query, collection_name], **kwargs)
        check_result = ResponseChecker(
            res, func_name, check_task, check_items, check, collection_name=collection_name, **kwargs
        ).run()
        return res, check_result

    @trace()
    def query_iterator(self, client, collection_name, timeout=None, check_task=None, check_items=None, **kwargs):
        timeout = TIMEOUT if timeout is None else timeout
        kwargs.update({"timeout": timeout})

        func_name = sys._getframe().f_code.co_name
        res, check = api_request([client.query_iterator, collection_name], **kwargs)
        check_result = ResponseChecker(
            res, func_name, check_task, check_items, check, collection_name=collection_name, **kwargs
        ).run()
        return res, check_result

    @trace()
    def get(
        self,
        client,
        collection_name,
        ids,
        output_fields=None,
        timeout=None,
        check_task=None,
        check_items=None,
        **kwargs,
    ):
        timeout = TIMEOUT if timeout is None else timeout
        kwargs.update({"timeout": timeout})

        func_name = sys._getframe().f_code.co_name
        res, check = api_request([client.get, collection_name, ids, output_fields], **kwargs)
        check_result = ResponseChecker(
            res,
            func_name,
            check_task,
            check_items,
            check,
            collection_name=collection_name,
            ids=ids,
            output_fields=output_fields,
            **kwargs,
        ).run()
        return res, check_result

    # No client.num_entities method
    # @trace()
    # def num_entities(self, client, collection_name, timeout=None, check_task=None, check_items=None, **kwargs):
    #     timeout = TIMEOUT if timeout is None else timeout
    #     kwargs.update({"timeout": timeout})

    #     func_name = sys._getframe().f_code.co_name
    #     res, check = api_request([client.num_entities, collection_name], **kwargs)
    #     check_result = ResponseChecker(res, func_name, check_task, check_items, check,
    #                                    collection_name=collection_name,
    #                                    **kwargs).run()
    #     return res, check_result

    @trace()
    def delete(
        self,
        client,
        collection_name,
        ids=None,
        timeout=None,
        filter=None,
        partition_name=None,
        check_task=None,
        check_items=None,
        **kwargs,
    ):
        timeout = TIMEOUT if timeout is None else timeout

        func_name = sys._getframe().f_code.co_name
        res, check = api_request([client.delete, collection_name, ids, timeout, filter, partition_name], **kwargs)
        check_result = ResponseChecker(
            res, func_name, check_task, check_items, check, collection_name=collection_name, **kwargs
        ).run()
        return res, check_result

    @trace()
    def flush(self, client, collection_name, timeout=None, check_task=None, check_items=None, **kwargs):
        timeout = TIMEOUT if timeout is None else timeout
        kwargs.update({"timeout": timeout})

        func_name = sys._getframe().f_code.co_name
        res, check = api_request([client.flush, collection_name], **kwargs)
        check_result = ResponseChecker(
            res, func_name, check_task, check_items, check, collection_name=collection_name, **kwargs
        ).run()
        return res, check_result

    @trace()
    def describe_collection(self, client, collection_name, timeout=None, check_task=None, check_items=None, **kwargs):
        timeout = TIMEOUT if timeout is None else timeout
        kwargs.update({"timeout": timeout})

        func_name = sys._getframe().f_code.co_name
        res, check = api_request([client.describe_collection, collection_name], **kwargs)
        check_result = ResponseChecker(
            res, func_name, check_task, check_items, check, collection_name=collection_name, **kwargs
        ).run()
        return res, check_result

    @trace()
    def list_collections(self, client, timeout=None, check_task=None, check_items=None, **kwargs):
        timeout = TIMEOUT if timeout is None else timeout
        kwargs.update({"timeout": timeout})

        func_name = sys._getframe().f_code.co_name
        res, check = api_request([client.list_collections], **kwargs)
        check_result = ResponseChecker(res, func_name, check_task, check_items, check, **kwargs).run()
        return res, check_result

    @trace()
    def drop_collection(self, client, collection_name, check_task=None, check_items=None, **kwargs):
        func_name = sys._getframe().f_code.co_name
        res, check = api_request([client.drop_collection, collection_name], **kwargs)
        check_result = ResponseChecker(
            res, func_name, check_task, check_items, check, collection_name=collection_name, **kwargs
        ).run()
        if check_result is True and collection_name in self.tear_down_collection_names:
            self.tear_down_collection_names.remove(collection_name)
        return res, check_result

    @trace()
    def list_indexes(self, client, collection_name, field_name=None, check_task=None, check_items=None, **kwargs):
        func_name = sys._getframe().f_code.co_name
        res, check = api_request([client.list_indexes, collection_name, field_name], **kwargs)
        check_result = ResponseChecker(
            res, func_name, check_task, check_items, check, collection_name=collection_name, **kwargs
        ).run()
        return res, check_result

    @trace()
    def describe_replica(self, client, collection_name, check_task=None, check_items=None, **kwargs):
        func_name = sys._getframe().f_code.co_name
        res, check = api_request([client.describe_replica, collection_name], **kwargs)
        check_result = ResponseChecker(
            res, func_name, check_task, check_items, check, collection_name=collection_name, **kwargs
        ).run()
        return res, check_result

    @trace()
    def get_load_state(self, client, collection_name, check_task=None, check_items=None, **kwargs):
        func_name = sys._getframe().f_code.co_name
        res, check = api_request([client.get_load_state, collection_name], **kwargs)
        check_result = ResponseChecker(
            res, func_name, check_task, check_items, check, collection_name=collection_name, **kwargs
        ).run()
        return res, check_result

    @trace()
    def load_collection(self, client, collection_name, timeout=None, check_task=None, check_items=None, **kwargs):
        timeout = TIMEOUT if timeout is None else timeout
        kwargs.update({"timeout": timeout})

        func_name = sys._getframe().f_code.co_name
        res, check = api_request([client.load_collection, collection_name], **kwargs)
        check_result = ResponseChecker(
            res, func_name, check_task, check_items, check, collection_name=collection_name, **kwargs
        ).run()
        return res, check_result

    @trace()
    def refresh_load(self, client, collection_name, timeout=None, check_task=None, check_items=None, **kwargs):
        timeout = TIMEOUT if timeout is None else timeout
        kwargs.update({"timeout": timeout})

        func_name = sys._getframe().f_code.co_name
        res, check = api_request([client.refresh_load, collection_name], **kwargs)
        check_result = ResponseChecker(
            res, func_name, check_task, check_items, check, collection_name=collection_name, **kwargs
        ).run()
        return res, check_result

    @trace()
    def release_collection(self, client, collection_name, timeout=None, check_task=None, check_items=None, **kwargs):
        timeout = TIMEOUT if timeout is None else timeout
        kwargs.update({"timeout": timeout})

        func_name = sys._getframe().f_code.co_name
        res, check = api_request([client.release_collection, collection_name], **kwargs)
        check_result = ResponseChecker(
            res, func_name, check_task, check_items, check, collection_name=collection_name, **kwargs
        ).run()
        return res, check_result

    @trace()
    def truncate_collection(self, client, collection_name, timeout=None, check_task=None, check_items=None, **kwargs):
        timeout = TIMEOUT if timeout is None else timeout
        kwargs.update({"timeout": timeout})

        func_name = sys._getframe().f_code.co_name
        res, check = api_request([client.truncate_collection, collection_name], **kwargs)
        check_result = ResponseChecker(
            res, func_name, check_task, check_items, check, collection_name=collection_name, **kwargs
        ).run()
        return res, check_result

    @trace()
    def list_persistent_segments(
        self, client, collection_name, timeout=None, check_task=None, check_items=None, **kwargs
    ):
        timeout = TIMEOUT if timeout is None else timeout
        kwargs.update({"timeout": timeout})

        func_name = sys._getframe().f_code.co_name
        res, check = api_request([client.list_persistent_segments, collection_name], **kwargs)
        check_result = ResponseChecker(
            res, func_name, check_task, check_items, check, collection_name=collection_name, **kwargs
        ).run()
        return res, check_result

    @trace()
    def load_partitions(
        self, client, collection_name, partition_names, timeout=None, check_task=None, check_items=None, **kwargs
    ):
        timeout = TIMEOUT if timeout is None else timeout
        kwargs.update({"timeout": timeout})

        func_name = sys._getframe().f_code.co_name
        res, check = api_request([client.load_partitions, collection_name, partition_names], **kwargs)
        check_result = ResponseChecker(
            res,
            func_name,
            check_task,
            check_items,
            check,
            collection_name=collection_name,
            partition_names=partition_names,
            **kwargs,
        ).run()
        return res, check_result

    @trace()
    def release_partitions(
        self, client, collection_name, partition_names, timeout=None, check_task=None, check_items=None, **kwargs
    ):
        timeout = TIMEOUT if timeout is None else timeout
        kwargs.update({"timeout": timeout})

        func_name = sys._getframe().f_code.co_name
        res, check = api_request([client.release_partitions, collection_name, partition_names], **kwargs)
        check_result = ResponseChecker(
            res,
            func_name,
            check_task,
            check_items,
            check,
            collection_name=collection_name,
            partition_names=partition_names,
            **kwargs,
        ).run()
        return res, check_result

    @trace()
    def rename_collection(self, client, old_name, new_name, timeout=None, check_task=None, check_items=None, **kwargs):
        timeout = TIMEOUT if timeout is None else timeout
        kwargs.update({"timeout": timeout})

        func_name = sys._getframe().f_code.co_name
        res, check = api_request([client.rename_collection, old_name, new_name], **kwargs)
        check_result = ResponseChecker(
            res, func_name, check_task, check_items, check, old_name=old_name, new_name=new_name, **kwargs
        ).run()
        self.tear_down_collection_names.append(new_name)
        return res, check_result

    @trace()
    def create_database(
        self, client, db_name, properties: dict | None = None, check_task=None, check_items=None, **kwargs
    ):
        func_name = sys._getframe().f_code.co_name
        res, check = api_request([client.create_database, db_name, properties], **kwargs)
        check_result = ResponseChecker(
            res, func_name, check_task, check_items, check, db_name=db_name, properties=properties, **kwargs
        ).run()
        return res, check_result

    @trace()
    def create_partition(
        self, client, collection_name, partition_name, timeout=None, check_task=None, check_items=None, **kwargs
    ):
        timeout = TIMEOUT if timeout is None else timeout
        kwargs.update({"timeout": timeout})

        func_name = sys._getframe().f_code.co_name
        res, check = api_request([client.create_partition, collection_name, partition_name], **kwargs)
        check_result = ResponseChecker(
            res,
            func_name,
            check_task,
            check_items,
            check,
            collection_name=collection_name,
            partition_name=partition_name,
            **kwargs,
        ).run()
        return res, check_result

    @trace()
    def list_partitions(self, client, collection_name, timeout=None, check_task=None, check_items=None, **kwargs):
        timeout = TIMEOUT if timeout is None else timeout
        kwargs.update({"timeout": timeout})

        func_name = sys._getframe().f_code.co_name
        res, check = api_request([client.list_partitions, collection_name], **kwargs)
        check_result = ResponseChecker(
            res, func_name, check_task, check_items, check, collection_name=collection_name, **kwargs
        ).run()
        return res, check_result

    @trace()
    def drop_partition(
        self, client, collection_name, partition_name, timeout=None, check_task=None, check_items=None, **kwargs
    ):
        timeout = TIMEOUT if timeout is None else timeout
        kwargs.update({"timeout": timeout})

        func_name = sys._getframe().f_code.co_name
        res, check = api_request([client.drop_partition, collection_name, partition_name], **kwargs)
        check_result = ResponseChecker(
            res,
            func_name,
            check_task,
            check_items,
            check,
            collection_name=collection_name,
            partition_name=partition_name,
            **kwargs,
        ).run()
        return res, check_result

    @trace()
    def has_partition(
        self, client, collection_name, partition_name, timeout=None, check_task=None, check_items=None, **kwargs
    ):
        timeout = TIMEOUT if timeout is None else timeout
        kwargs.update({"timeout": timeout})

        func_name = sys._getframe().f_code.co_name
        res, check = api_request([client.has_partition, collection_name, partition_name], **kwargs)
        check_result = ResponseChecker(
            res,
            func_name,
            check_task,
            check_items,
            check,
            collection_name=collection_name,
            partition_name=partition_name,
            **kwargs,
        ).run()
        return res, check_result

    @trace()
    def get_partition_stats(
        self, client, collection_name, partition_name, timeout=None, check_task=None, check_items=None, **kwargs
    ):
        timeout = TIMEOUT if timeout is None else timeout
        kwargs.update({"timeout": timeout})

        func_name = sys._getframe().f_code.co_name
        res, check = api_request([client.get_partition_stats, collection_name, partition_name], **kwargs)
        check_result = ResponseChecker(
            res,
            func_name,
            check_task,
            check_items,
            check,
            collection_name=collection_name,
            partition_name=partition_name,
            **kwargs,
        ).run()
        return res, check_result

    @trace()
    def prepare_index_params(self, client, check_task=None, check_items=None, **kwargs):

        func_name = sys._getframe().f_code.co_name
        res, check = api_request([client.prepare_index_params], **kwargs)
        check_result = ResponseChecker(res, func_name, check_task, check_items, check, **kwargs).run()
        return res, check_result

    @trace()
    def create_index(
        self, client, collection_name, index_params, timeout=None, check_task=None, check_items=None, **kwargs
    ):
        timeout = TIMEOUT if timeout is None else timeout
        kwargs.update({"timeout": timeout})

        func_name = sys._getframe().f_code.co_name
        res, check = api_request([client.create_index, collection_name, index_params], **kwargs)
        check_result = ResponseChecker(
            res,
            func_name,
            check_task,
            check_items,
            check,
            collection_name=collection_name,
            index_params=index_params,
            **kwargs,
        ).run()
        return res, check_result

    @trace()
    def drop_index(
        self, client, collection_name, index_name, timeout=None, check_task=None, check_items=None, **kwargs
    ):
        timeout = TIMEOUT if timeout is None else timeout
        kwargs.update({"timeout": timeout})

        func_name = sys._getframe().f_code.co_name
        res, check = api_request([client.drop_index, collection_name, index_name], **kwargs)
        check_result = ResponseChecker(
            res,
            func_name,
            check_task,
            check_items,
            check,
            collection_name=collection_name,
            index_name=index_name,
            **kwargs,
        ).run()
        return res, check_result

    @trace()
    def describe_index(
        self, client, collection_name, index_name, timeout=None, check_task=None, check_items=None, **kwargs
    ):
        timeout = TIMEOUT if timeout is None else timeout
        kwargs.update({"timeout": timeout})

        func_name = sys._getframe().f_code.co_name
        res, check = api_request([client.describe_index, collection_name, index_name], **kwargs)
        check_result = ResponseChecker(
            res,
            func_name,
            check_task,
            check_items,
            check,
            collection_name=collection_name,
            index_name=index_name,
            **kwargs,
        ).run()
        return res, check_result

    def wait_for_index_ready(self, client, collection_name, index_name, timeout=None, **kwargs):
        timeout = TIMEOUT if timeout is None else timeout
        start_time = time.time()
        while start_time + timeout > time.time():
            index_info, _ = self.describe_index(client, collection_name, index_name, **kwargs)
            if index_info.get("pending_index_rows", 1) == 0:
                return True
            time.sleep(2)
        return False

    def wait_for_compaction_ready(self, client, compact_id, timeout=None, **kwargs):
        timeout = TIMEOUT if timeout is None else timeout
        start_time = time.time()
        while start_time + timeout > time.time():
            res = self.get_compaction_state(client, compact_id, **kwargs)[0]
            if res == "Completed":
                return True
            time.sleep(2)
        return False

    @trace()
    def create_alias(self, client, collection_name, alias, timeout=None, check_task=None, check_items=None, **kwargs):
        timeout = TIMEOUT if timeout is None else timeout
        kwargs.update({"timeout": timeout})

        func_name = sys._getframe().f_code.co_name
        res, check = api_request([client.create_alias, collection_name, alias], **kwargs)
        check_result = ResponseChecker(
            res, func_name, check_task, check_items, check, collection_name=collection_name, alias=alias, **kwargs
        ).run()
        return res, check_result

    @trace()
    def drop_alias(self, client, alias, timeout=None, check_task=None, check_items=None, **kwargs):
        timeout = TIMEOUT if timeout is None else timeout
        kwargs.update({"timeout": timeout})

        func_name = sys._getframe().f_code.co_name
        res, check = api_request([client.drop_alias, alias], **kwargs)
        check_result = ResponseChecker(res, func_name, check_task, check_items, check, alias=alias, **kwargs).run()
        return res, check_result

    @trace()
    def alter_alias(self, client, collection_name, alias, timeout=None, check_task=None, check_items=None, **kwargs):
        timeout = TIMEOUT if timeout is None else timeout
        kwargs.update({"timeout": timeout})

        func_name = sys._getframe().f_code.co_name
        res, check = api_request([client.alter_alias, collection_name, alias], **kwargs)
        check_result = ResponseChecker(
            res, func_name, check_task, check_items, check, collection_name=collection_name, alias=alias, **kwargs
        ).run()
        return res, check_result

    @trace()
    def describe_alias(self, client, alias, timeout=None, check_task=None, check_items=None, **kwargs):
        timeout = TIMEOUT if timeout is None else timeout
        kwargs.update({"timeout": timeout})

        func_name = sys._getframe().f_code.co_name
        res, check = api_request([client.describe_alias, alias], **kwargs)
        check_result = ResponseChecker(res, func_name, check_task, check_items, check, alias=alias, **kwargs).run()
        return res, check_result

    @trace()
    def list_aliases(self, client, collection_name, timeout=None, check_task=None, check_items=None, **kwargs):
        timeout = TIMEOUT if timeout is None else timeout
        kwargs.update({"timeout": timeout})

        func_name = sys._getframe().f_code.co_name
        res, check = api_request([client.list_aliases, collection_name], **kwargs)
        check_result = ResponseChecker(
            res, func_name, check_task, check_items, check, collection_name=collection_name, **kwargs
        ).run()
        return res, check_result

    @trace()
    def using_database(self, client, db_name, timeout=None, check_task=None, check_items=None, **kwargs):
        timeout = TIMEOUT if timeout is None else timeout
        kwargs.update({"timeout": timeout})

        func_name = sys._getframe().f_code.co_name
        res, check = api_request([client.using_database, db_name], **kwargs)
        check_result = ResponseChecker(res, func_name, check_task, check_items, check, **kwargs).run()
        return res, check_result

    @trace()
    def use_database(self, client, db_name, timeout=None, check_task=None, check_items=None, **kwargs):
        timeout = TIMEOUT if timeout is None else timeout
        kwargs.update({"timeout": timeout})

        func_name = sys._getframe().f_code.co_name
        res, check = api_request([client.use_database, db_name], **kwargs)
        check_result = ResponseChecker(res, func_name, check_task, check_items, check, **kwargs).run()
        return res, check_result

    def create_user(self, client, user_name, password, timeout=None, check_task=None, check_items=None, **kwargs):
        timeout = TIMEOUT if timeout is None else timeout
        kwargs.update({"timeout": timeout})
        func_name = sys._getframe().f_code.co_name
        res, check = api_request([client.create_user, user_name, password], **kwargs)
        check_result = ResponseChecker(
            res, func_name, check_task, check_items, check, user_name=user_name, password=password, **kwargs
        ).run()
        return res, check_result

    @trace()
    def drop_user(self, client, user_name, timeout=None, check_task=None, check_items=None, **kwargs):
        timeout = TIMEOUT if timeout is None else timeout
        kwargs.update({"timeout": timeout})
        func_name = sys._getframe().f_code.co_name
        res, check = api_request([client.drop_user, user_name], **kwargs)
        check_result = ResponseChecker(
            res, func_name, check_task, check_items, check, user_name=user_name, **kwargs
        ).run()
        return res, check_result

    @trace()
    def update_password(
        self,
        client,
        user_name,
        old_password,
        new_password,
        reset_connection=False,
        timeout=None,
        check_task=None,
        check_items=None,
        **kwargs,
    ):
        timeout = TIMEOUT if timeout is None else timeout
        kwargs.update({"timeout": timeout})
        func_name = sys._getframe().f_code.co_name
        res, check = api_request(
            [client.update_password, user_name, old_password, new_password, reset_connection], **kwargs
        )
        check_result = ResponseChecker(
            res,
            func_name,
            check_task,
            check_items,
            check,
            user_name=user_name,
            old_password=old_password,
            new_password=new_password,
            reset_connection=reset_connection,
            **kwargs,
        ).run()
        return res, check_result

    @trace()
    def list_users(self, client, timeout=None, check_task=None, check_items=None, **kwargs):
        timeout = TIMEOUT if timeout is None else timeout
        kwargs.update({"timeout": timeout})
        func_name = sys._getframe().f_code.co_name
        res, check = api_request([client.list_users], **kwargs)
        check_result = ResponseChecker(res, func_name, check_task, check_items, check, **kwargs).run()
        return res, check_result

    @trace()
    def describe_user(self, client, user_name, timeout=None, check_task=None, check_items=None, **kwargs):
        timeout = TIMEOUT if timeout is None else timeout
        kwargs.update({"timeout": timeout})
        func_name = sys._getframe().f_code.co_name
        res, check = api_request([client.describe_user, user_name], **kwargs)
        check_result = ResponseChecker(
            res, func_name, check_task, check_items, check, user_name=user_name, **kwargs
        ).run()
        return res, check_result

    @trace()
    def create_role(self, client, role_name, timeout=None, check_task=None, check_items=None, **kwargs):
        timeout = TIMEOUT if timeout is None else timeout
        kwargs.update({"timeout": timeout})
        func_name = sys._getframe().f_code.co_name
        res, check = api_request([client.create_role, role_name], **kwargs)
        check_result = ResponseChecker(
            res, func_name, check_task, check_items, check, role_name=role_name, **kwargs
        ).run()
        return res, check_result

    @trace()
    def drop_role(self, client, role_name, timeout=None, check_task=None, check_items=None, **kwargs):
        timeout = TIMEOUT if timeout is None else timeout
        kwargs.update({"timeout": timeout})
        func_name = sys._getframe().f_code.co_name
        res, check = api_request([client.drop_role, role_name], **kwargs)
        check_result = ResponseChecker(
            res, func_name, check_task, check_items, check, role_name=role_name, **kwargs
        ).run()
        return res, check_result

    @trace()
    def describe_role(self, client, role_name, timeout=None, check_task=None, check_items=None, **kwargs):
        timeout = TIMEOUT if timeout is None else timeout
        kwargs.update({"timeout": timeout})
        func_name = sys._getframe().f_code.co_name
        res, check = api_request([client.describe_role, role_name], **kwargs)
        check_result = ResponseChecker(
            res, func_name, check_task, check_items, check, role_name=role_name, **kwargs
        ).run()
        return res, check_result

    @trace()
    def list_roles(self, client, timeout=None, check_task=None, check_items=None, **kwargs):
        timeout = TIMEOUT if timeout is None else timeout
        kwargs.update({"timeout": timeout})
        func_name = sys._getframe().f_code.co_name
        res, check = api_request([client.list_roles], **kwargs)
        check_result = ResponseChecker(res, func_name, check_task, check_items, check, **kwargs).run()
        return res, check_result

    @trace()
    def grant_role(self, client, user_name, role_name, timeout=None, check_task=None, check_items=None, **kwargs):
        timeout = TIMEOUT if timeout is None else timeout
        kwargs.update({"timeout": timeout})
        func_name = sys._getframe().f_code.co_name
        res, check = api_request([client.grant_role, user_name, role_name], **kwargs)
        check_result = ResponseChecker(
            res, func_name, check_task, check_items, check, user_name=user_name, role_name=role_name, **kwargs
        ).run()
        return res, check_result

    @trace()
    def revoke_role(self, client, user_name, role_name, timeout=None, check_task=None, check_items=None, **kwargs):
        timeout = TIMEOUT if timeout is None else timeout
        kwargs.update({"timeout": timeout})
        func_name = sys._getframe().f_code.co_name
        res, check = api_request([client.revoke_role, user_name, role_name], **kwargs)
        check_result = ResponseChecker(
            res, func_name, check_task, check_items, check, user_name=user_name, role_name=role_name, **kwargs
        ).run()
        return res, check_result

    @trace()
    def grant_privilege(
        self,
        client,
        role_name,
        object_type,
        privilege,
        object_name,
        db_name="",
        timeout=None,
        check_task=None,
        check_items=None,
        **kwargs,
    ):
        timeout = TIMEOUT if timeout is None else timeout
        kwargs.update({"timeout": timeout})
        func_name = sys._getframe().f_code.co_name
        res, check = api_request(
            [client.grant_privilege, role_name, object_type, privilege, object_name, db_name], **kwargs
        )
        check_result = ResponseChecker(
            res,
            func_name,
            check_task,
            check_items,
            check,
            role_name=role_name,
            object_type=object_type,
            privilege=privilege,
            object_name=object_name,
            db_name=db_name,
            **kwargs,
        ).run()
        return res, check_result

    @trace()
    def revoke_privilege(
        self,
        client,
        role_name,
        object_type,
        privilege,
        object_name,
        db_name="",
        timeout=None,
        check_task=None,
        check_items=None,
        **kwargs,
    ):
        timeout = TIMEOUT if timeout is None else timeout
        kwargs.update({"timeout": timeout})
        func_name = sys._getframe().f_code.co_name
        res, check = api_request(
            [client.revoke_privilege, role_name, object_type, privilege, object_name, db_name], **kwargs
        )
        check_result = ResponseChecker(
            res,
            func_name,
            check_task,
            check_items,
            check,
            role_name=role_name,
            object_type=object_type,
            privilege=privilege,
            object_name=object_name,
            db_name=db_name,
            **kwargs,
        ).run()
        return res, check_result

    @trace()
    def create_privilege_group(
        self, client, privilege_group: str, timeout=None, check_task=None, check_items=None, **kwargs
    ):
        timeout = TIMEOUT if timeout is None else timeout
        kwargs.update({"timeout": timeout})

        func_name = sys._getframe().f_code.co_name
        res, check = api_request([client.create_privilege_group, privilege_group], **kwargs)
        check_result = ResponseChecker(
            res, func_name, check_task, check_items, check, privilege_group=privilege_group, **kwargs
        ).run()
        return res, check_result

    @trace()
    def drop_privilege_group(
        self, client, privilege_group: str, timeout=None, check_task=None, check_items=None, **kwargs
    ):
        timeout = TIMEOUT if timeout is None else timeout
        kwargs.update({"timeout": timeout})

        func_name = sys._getframe().f_code.co_name
        res, check = api_request([client.drop_privilege_group, privilege_group], **kwargs)
        check_result = ResponseChecker(
            res, func_name, check_task, check_items, check, privilege_group=privilege_group, **kwargs
        ).run()
        return res, check_result

    @trace()
    def list_privilege_groups(self, client, timeout=None, check_task=None, check_items=None, **kwargs):
        timeout = TIMEOUT if timeout is None else timeout
        kwargs.update({"timeout": timeout})

        func_name = sys._getframe().f_code.co_name
        res, check = api_request([client.list_privilege_groups], **kwargs)
        check_result = ResponseChecker(res, func_name, check_task, check_items, check, **kwargs).run()
        return res, check_result

    @trace()
    def add_privileges_to_group(
        self, client, privilege_group: str, privileges: list, timeout=None, check_task=None, check_items=None, **kwargs
    ):
        timeout = TIMEOUT if timeout is None else timeout
        kwargs.update({"timeout": timeout})

        func_name = sys._getframe().f_code.co_name
        res, check = api_request([client.add_privileges_to_group, privilege_group, privileges], **kwargs)
        check_result = ResponseChecker(
            res,
            func_name,
            check_task,
            check_items,
            check,
            privilege_group=privilege_group,
            privileges=privileges,
            **kwargs,
        ).run()
        return res, check_result

    @trace()
    def remove_privileges_from_group(
        self, client, privilege_group: str, privileges: list, timeout=None, check_task=None, check_items=None, **kwargs
    ):
        timeout = TIMEOUT if timeout is None else timeout
        kwargs.update({"timeout": timeout})

        func_name = sys._getframe().f_code.co_name
        res, check = api_request([client.remove_privileges_from_group, privilege_group, privileges], **kwargs)
        check_result = ResponseChecker(
            res,
            func_name,
            check_task,
            check_items,
            check,
            privilege_group=privilege_group,
            privileges=privileges,
            **kwargs,
        ).run()
        return res, check_result

    @trace()
    def grant_privilege_v2(
        self,
        client,
        role_name: str,
        privilege: str,
        collection_name: str,
        db_name=None,
        timeout=None,
        check_task=None,
        check_items=None,
        **kwargs,
    ):
        timeout = TIMEOUT if timeout is None else timeout
        kwargs.update({"timeout": timeout})

        func_name = sys._getframe().f_code.co_name
        res, check = api_request([client.grant_privilege_v2, role_name, privilege, collection_name, db_name], **kwargs)
        check_result = ResponseChecker(
            res,
            func_name,
            check_task,
            check_items,
            check,
            role_name=role_name,
            privilege=privilege,
            collection_name=collection_name,
            db_name=db_name,
            **kwargs,
        ).run()
        return res, check_result

    @trace()
    def revoke_privilege_v2(
        self,
        client,
        role_name: str,
        privilege: str,
        collection_name: str,
        db_name=None,
        timeout=None,
        check_task=None,
        check_items=None,
        **kwargs,
    ):
        timeout = TIMEOUT if timeout is None else timeout
        kwargs.update({"timeout": timeout})

        func_name = sys._getframe().f_code.co_name
        res, check = api_request([client.revoke_privilege_v2, role_name, privilege, collection_name, db_name], **kwargs)
        check_result = ResponseChecker(
            res,
            func_name,
            check_task,
            check_items,
            check,
            role_name=role_name,
            privilege=privilege,
            collection_name=collection_name,
            db_name=db_name,
            **kwargs,
        ).run()
        return res, check_result

    @trace()
    def alter_index_properties(
        self, client, collection_name, index_name, properties, timeout=None, check_task=None, check_items=None, **kwargs
    ):
        timeout = TIMEOUT if timeout is None else timeout
        kwargs.update({"timeout": timeout})

        func_name = sys._getframe().f_code.co_name
        res, check = api_request([client.alter_index_properties, collection_name, index_name, properties], **kwargs)
        check_result = ResponseChecker(res, func_name, check_task, check_items, check, **kwargs).run()
        return res, check_result

    @trace()
    def drop_index_properties(
        self,
        client,
        collection_name,
        index_name,
        property_keys,
        timeout=None,
        check_task=None,
        check_items=None,
        **kwargs,
    ):
        timeout = TIMEOUT if timeout is None else timeout
        kwargs.update({"timeout": timeout})

        func_name = sys._getframe().f_code.co_name
        res, check = api_request([client.drop_index_properties, collection_name, index_name, property_keys], **kwargs)
        check_result = ResponseChecker(res, func_name, check_task, check_items, check, **kwargs).run()
        return res, check_result

    @trace()
    def alter_collection_properties(
        self, client, collection_name, properties, timeout=None, check_task=None, check_items=None, **kwargs
    ):
        timeout = TIMEOUT if timeout is None else timeout
        kwargs.update({"timeout": timeout})

        func_name = sys._getframe().f_code.co_name
        res, check = api_request([client.alter_collection_properties, collection_name, properties], **kwargs)
        check_result = ResponseChecker(res, func_name, check_task, check_items, check, **kwargs).run()
        return res, check_result

    @trace()
    def drop_collection_properties(
        self, client, collection_name, property_keys, timeout=None, check_task=None, check_items=None, **kwargs
    ):
        timeout = TIMEOUT if timeout is None else timeout

        func_name = sys._getframe().f_code.co_name
        res, check = api_request([client.drop_collection_properties, collection_name, property_keys, timeout], **kwargs)
        check_result = ResponseChecker(res, func_name, check_task, check_items, check, **kwargs).run()
        return res, check_result

    @trace()
    def alter_collection_field(
        self,
        client,
        collection_name,
        field_name,
        field_params,
        timeout=None,
        check_task=None,
        check_items=None,
        **kwargs,
    ):
        timeout = TIMEOUT if timeout is None else timeout

        func_name = sys._getframe().f_code.co_name
        res, check = api_request(
            [client.alter_collection_field, collection_name, field_name, field_params, timeout], **kwargs
        )
        check_result = ResponseChecker(res, func_name, check_task, check_items, check, **kwargs).run()
        return res, check_result

    @trace()
    def alter_database_properties(
        self, client, db_name, properties, timeout=None, check_task=None, check_items=None, **kwargs
    ):
        timeout = TIMEOUT if timeout is None else timeout
        kwargs.update({"timeout": timeout})

        func_name = sys._getframe().f_code.co_name
        res, check = api_request([client.alter_database_properties, db_name, properties], **kwargs)
        check_result = ResponseChecker(res, func_name, check_task, check_items, check, **kwargs).run()
        return res, check_result

    @trace()
    def drop_database_properties(
        self, client, db_name, property_keys, timeout=None, check_task=None, check_items=None, **kwargs
    ):
        timeout = TIMEOUT if timeout is None else timeout
        kwargs.update({"timeout": timeout})

        func_name = sys._getframe().f_code.co_name
        res, check = api_request([client.drop_database_properties, db_name, property_keys], **kwargs)
        check_result = ResponseChecker(res, func_name, check_task, check_items, check, **kwargs).run()
        return res, check_result

    @trace()
    def describe_database(self, client, db_name, timeout=None, check_task=None, check_items=None, **kwargs):
        timeout = TIMEOUT if timeout is None else timeout
        kwargs.update({"timeout": timeout})

        func_name = sys._getframe().f_code.co_name
        res, check = api_request([client.describe_database, db_name], **kwargs)
        check_result = ResponseChecker(res, func_name, check_task, check_items, check, **kwargs).run()
        return res, check_result

    @trace()
    def drop_database(self, client, db_name, timeout=None, check_task=None, check_items=None, **kwargs):
        timeout = TIMEOUT if timeout is None else timeout
        kwargs.update({"timeout": timeout})

        func_name = sys._getframe().f_code.co_name
        res, check = api_request([client.drop_database, db_name], **kwargs)
        check_result = ResponseChecker(res, func_name, check_task, check_items, check, **kwargs).run()
        return res, check_result

    @trace()
    def list_databases(self, client, timeout=None, check_task=None, check_items=None, **kwargs):
        timeout = TIMEOUT if timeout is None else timeout
        kwargs.update({"timeout": timeout})

        func_name = sys._getframe().f_code.co_name
        res, check = api_request([client.list_databases], **kwargs)
        check_result = ResponseChecker(res, func_name, check_task, check_items, check, **kwargs).run()
        return res, check_result

    @trace()
    def run_analyzer(self, client, text, analyzer_params, timeout=None, check_task=None, check_items=None, **kwargs):
        timeout = TIMEOUT if timeout is None else timeout
        kwargs.update({"timeout": timeout})
        func_name = sys._getframe().f_code.co_name
        res, check = api_request([client.run_analyzer, text, analyzer_params], **kwargs)
        check_result = ResponseChecker(
            res, func_name, check_task, check_items, check, text=text, analyzer_params=analyzer_params, **kwargs
        ).run()
        return res, check_result

    def compact(
        self, client, collection_name, is_clustering=False, timeout=None, check_task=None, check_items=None, **kwargs
    ):
        timeout = TIMEOUT if timeout is None else timeout
        kwargs.update({"timeout": timeout})

        func_name = sys._getframe().f_code.co_name
        res, check = api_request([client.compact, collection_name, is_clustering], **kwargs)
        check_result = ResponseChecker(
            res,
            func_name,
            check_task,
            check_items,
            check,
            collection_name=collection_name,
            is_clustering=is_clustering,
            **kwargs,
        ).run()
        return res, check_result

    @trace()
    def get_compaction_state(self, client, job_id, timeout=None, check_task=None, check_items=None, **kwargs):
        timeout = TIMEOUT if timeout is None else timeout
        kwargs.update({"timeout": timeout})

        func_name = sys._getframe().f_code.co_name
        res, check = api_request([client.get_compaction_state, job_id], **kwargs)
        check_result = ResponseChecker(res, func_name, check_task, check_items, check, job_id=job_id, **kwargs).run()
        return res, check_result

    @trace()
    def create_resource_group(self, client, name, timeout=None, check_task=None, check_items=None, **kwargs):
        timeout = TIMEOUT if timeout is None else timeout
        kwargs.update({"timeout": timeout})

        func_name = sys._getframe().f_code.co_name
        res, check = api_request([client.create_resource_group, name], **kwargs)
        check_result = ResponseChecker(res, func_name, check_task, check_items, check, name=name, **kwargs).run()
        return res, check_result

    @trace()
    def update_resource_groups(self, client, configs, timeout=None, check_task=None, check_items=None, **kwargs):
        timeout = TIMEOUT if timeout is None else timeout
        kwargs.update({"timeout": timeout})

        func_name = sys._getframe().f_code.co_name
        res, check = api_request([client.update_resource_groups, configs], **kwargs)
        check_result = ResponseChecker(res, func_name, check_task, check_items, check, configs=configs, **kwargs).run()
        return res, check_result

    @trace()
    def drop_resource_group(self, client, name, timeout=None, check_task=None, check_items=None, **kwargs):
        timeout = TIMEOUT if timeout is None else timeout
        kwargs.update({"timeout": timeout})

        func_name = sys._getframe().f_code.co_name
        res, check = api_request([client.drop_resource_group, name], **kwargs)
        check_result = ResponseChecker(res, func_name, check_task, check_items, check, name=name, **kwargs).run()
        return res, check_result

    @trace()
    def describe_resource_group(self, client, name, timeout=None, check_task=None, check_items=None, **kwargs):
        timeout = TIMEOUT if timeout is None else timeout
        kwargs.update({"timeout": timeout})

        func_name = sys._getframe().f_code.co_name
        res, check = api_request([client.describe_resource_group, name], **kwargs)
        check_result = ResponseChecker(res, func_name, check_task, check_items, check, name=name, **kwargs).run()
        return res, check_result

    @trace()
    def list_resource_groups(self, client, timeout=None, check_task=None, check_items=None, **kwargs):
        timeout = TIMEOUT if timeout is None else timeout
        kwargs.update({"timeout": timeout})

        func_name = sys._getframe().f_code.co_name
        res, check = api_request([client.list_resource_groups], **kwargs)
        check_result = ResponseChecker(res, func_name, check_task, check_items, check, **kwargs).run()
        return res, check_result

    @trace()
    def transfer_replica(
        self,
        client,
        source_group,
        target_group,
        collection_name,
        num_replicas,
        timeout=None,
        check_task=None,
        check_items=None,
        **kwargs,
    ):
        timeout = TIMEOUT if timeout is None else timeout
        kwargs.update({"timeout": timeout})

        func_name = sys._getframe().f_code.co_name
        res, check = api_request(
            [client.transfer_replica, source_group, target_group, collection_name, num_replicas], **kwargs
        )
        check_result = ResponseChecker(
            res,
            func_name,
            check_task,
            check_items,
            check,
            source_group=source_group,
            target_group=target_group,
            collection_name=collection_name,
            num_replicas=num_replicas,
            **kwargs,
        ).run()
        return res, check_result

    @trace()
    def add_collection_field(
        self,
        client,
        collection_name,
        field_name,
        data_type,
        desc="",
        timeout=None,
        check_task=None,
        check_items=None,
        **kwargs,
    ):
        timeout = TIMEOUT if timeout is None else timeout
        kwargs.update({"timeout": timeout})

        func_name = sys._getframe().f_code.co_name
        res, check = api_request([client.add_collection_field, collection_name, field_name, data_type, desc], **kwargs)
        check_result = ResponseChecker(res, func_name, check_task, check_items, check, **kwargs).run()
        return res, check_result

    def wait_schema_version_consistent(self, client, collection_name, timeout=30, poll_interval=0.01):
        """
        Poll get_collection_stats until the schema version consistency gate would pass.

        Background: After a previous schema-change DDL (AlterCollectionSchema or
        AddCollectionField), DataCoord's backfill segment-version propagation runs on
        a periodic tick. Until the tick fires, a back-to-back schema-change call is
        rejected by the consistency gate at the Proxy / RootCoord with an error like
        "schema version consistency check failed: N/M segments have caught up".
        E2E tests that issue successive schema changes need to wait until the gate
        would pass; this helper polls the same stats keys the gate reads.

        Pass condition (mirrors checkSchemaVersionConsistency):
          - both keys absent → schema version is 0, no backfill in progress → pass
          - schema_version_consistent_segments == schema_version_total_segments → pass

        Note: pymilvus MilvusClient.get_collection_stats returns a dict[str, str]
        (with row_count converted to int), so we read the keys directly from the dict.

        Args:
            client: pymilvus MilvusClient
            collection_name: target collection
            timeout: max seconds to wait before giving up (default 30s)
            poll_interval: sleep between polls in seconds (default 10ms)

        Raises:
            AssertionError on timeout, with the last observed counts for debugging.
        """
        consistent_key = "schema_version_consistent_segments"
        total_key = "schema_version_total_segments"
        deadline = time.time() + timeout
        last_consistent, last_total = None, None
        while time.time() < deadline:
            stats, _ = self.get_collection_stats(client, collection_name)
            if not stats:
                time.sleep(poll_interval)
                continue
            consistent = stats.get(consistent_key)
            total = stats.get(total_key)
            # Both absent → schema v0 path, gate passes trivially.
            if consistent is None and total is None:
                return
            # Both present and equal → backfill caught up.
            if consistent is not None and total is not None and int(consistent) == int(total):
                return
            last_consistent, last_total = consistent, total
            time.sleep(poll_interval)
        raise AssertionError(
            f"wait_schema_version_consistent timed out after {timeout}s for collection "
            f"{collection_name}: consistent={last_consistent}, total={last_total}"
        )

    def add_collection_field_wait_schema_version_consistency(
        self,
        client,
        collection_name,
        field_name,
        data_type,
        desc="",
        timeout=None,
        check_task=None,
        check_items=None,
        wait_timeout=30,
        poll_interval=0.01,
        **kwargs,
    ):
        """
        Wrapper around add_collection_field that first waits for the schema-version
        consistency gate to pass. Use this in E2E tests that issue successive
        schema-change DDLs back-to-back, where the previous call's backfill
        segment-version propagation tick may not have fired yet.

        See wait_schema_version_consistent for the polling logic and pass conditions.
        All add_collection_field arguments are forwarded unchanged.
        """
        self.wait_schema_version_consistent(client, collection_name, timeout=wait_timeout, poll_interval=poll_interval)
        return self.add_collection_field(
            client,
            collection_name,
            field_name,
            data_type,
            desc=desc,
            timeout=timeout,
            check_task=check_task,
            check_items=check_items,
            **kwargs,
        )

    # ====================== Snapshot ======================
    @trace()
    def create_snapshot(
        self,
        client,
        collection_name,
        snapshot_name,
        description="",
        timeout=None,
        check_task=None,
        check_items=None,
        **kwargs,
    ):
        """Create a snapshot for a collection.

        Note: wrapper keeps ``collection_name`` before ``snapshot_name`` for test
        readability, but the SDK signature is ``create_snapshot(snapshot_name,
        collection_name, ...)`` so we forward positionals in SDK order.
        """
        timeout = TIMEOUT if timeout is None else timeout
        kwargs.update({"timeout": timeout})

        func_name = sys._getframe().f_code.co_name
        res, check = api_request(
            [client.create_snapshot, snapshot_name, collection_name], description=description, **kwargs
        )
        check_result = ResponseChecker(
            res,
            func_name,
            check_task,
            check_items,
            check,
            snapshot_name=snapshot_name,
            collection_name=collection_name,
            **kwargs,
        ).run()
        return res, check_result

    @trace()
    def drop_snapshot(
        self, client, snapshot_name, collection_name, timeout=None, check_task=None, check_items=None, **kwargs
    ):
        timeout = TIMEOUT if timeout is None else timeout
        kwargs.update({"timeout": timeout})

        func_name = sys._getframe().f_code.co_name
        res, check = api_request([client.drop_snapshot, snapshot_name, collection_name], **kwargs)
        check_result = ResponseChecker(
            res,
            func_name,
            check_task,
            check_items,
            check,
            snapshot_name=snapshot_name,
            collection_name=collection_name,
            **kwargs,
        ).run()
        return res, check_result

    @trace()
    def list_snapshots(self, client, collection_name="", timeout=None, check_task=None, check_items=None, **kwargs):
        timeout = TIMEOUT if timeout is None else timeout
        kwargs.update({"timeout": timeout})

        func_name = sys._getframe().f_code.co_name
        res, check = api_request([client.list_snapshots, collection_name], **kwargs)
        check_result = ResponseChecker(
            res, func_name, check_task, check_items, check, collection_name=collection_name, **kwargs
        ).run()
        return res, check_result

    @trace()
    def describe_snapshot(
        self, client, snapshot_name, collection_name, timeout=None, check_task=None, check_items=None, **kwargs
    ):
        timeout = TIMEOUT if timeout is None else timeout
        kwargs.update({"timeout": timeout})

        func_name = sys._getframe().f_code.co_name
        res, check = api_request([client.describe_snapshot, snapshot_name, collection_name], **kwargs)
        check_result = ResponseChecker(
            res,
            func_name,
            check_task,
            check_items,
            check,
            snapshot_name=snapshot_name,
            collection_name=collection_name,
            **kwargs,
        ).run()
        return res, check_result

    @trace()
    def restore_snapshot(
        self,
        client,
        snapshot_name,
        target_collection_name,
        source_collection_name="",
        timeout=None,
        check_task=None,
        check_items=None,
        **kwargs,
    ):
        """Restore a snapshot into a new collection.

        SDK positional order is ``(snapshot_name, source_collection_name,
        target_collection_name)`` and ``source_collection_name`` is now required.
        """
        timeout = TIMEOUT if timeout is None else timeout
        kwargs.update({"timeout": timeout})

        func_name = sys._getframe().f_code.co_name
        res, check = api_request(
            [client.restore_snapshot, snapshot_name, source_collection_name, target_collection_name], **kwargs
        )
        check_result = ResponseChecker(
            res,
            func_name,
            check_task,
            check_items,
            check,
            snapshot_name=snapshot_name,
            target_collection_name=target_collection_name,
            source_collection_name=source_collection_name,
            **kwargs,
        ).run()
        return res, check_result

    @trace()
    def get_restore_snapshot_state(self, client, job_id, timeout=None, check_task=None, check_items=None, **kwargs):
        timeout = TIMEOUT if timeout is None else timeout
        kwargs.update({"timeout": timeout})

        func_name = sys._getframe().f_code.co_name
        res, check = api_request([client.get_restore_snapshot_state, job_id], **kwargs)
        check_result = ResponseChecker(res, func_name, check_task, check_items, check, job_id=job_id, **kwargs).run()
        return res, check_result

    @trace()
    def list_restore_snapshot_jobs(
        self, client, collection_name="", timeout=None, check_task=None, check_items=None, **kwargs
    ):
        timeout = TIMEOUT if timeout is None else timeout
        kwargs.update({"timeout": timeout})

        func_name = sys._getframe().f_code.co_name
        res, check = api_request([client.list_restore_snapshot_jobs, collection_name], **kwargs)
        check_result = ResponseChecker(
            res, func_name, check_task, check_items, check, collection_name=collection_name, **kwargs
        ).run()
        return res, check_result

    @trace()
    def pin_snapshot_data(
        self,
        client,
        snapshot_name,
        collection_name,
        ttl_seconds=0,
        timeout=None,
        check_task=None,
        check_items=None,
        **kwargs,
    ):
        """Pin snapshot data to prevent GC / compaction. Returns pin_id (int).

        SDK signature: ``pin_snapshot_data(snapshot_name, collection_name,
        db_name="", ttl_seconds=0, ...)``.
        """
        timeout = TIMEOUT if timeout is None else timeout
        kwargs.update({"timeout": timeout})

        func_name = sys._getframe().f_code.co_name
        res, check = api_request(
            [client.pin_snapshot_data, snapshot_name, collection_name], ttl_seconds=ttl_seconds, **kwargs
        )
        check_result = ResponseChecker(
            res,
            func_name,
            check_task,
            check_items,
            check,
            snapshot_name=snapshot_name,
            collection_name=collection_name,
            ttl_seconds=ttl_seconds,
            **kwargs,
        ).run()
        return res, check_result

    @trace()
    def unpin_snapshot_data(self, client, pin_id, timeout=None, check_task=None, check_items=None, **kwargs):
        """Release a pin previously created by ``pin_snapshot_data``."""
        timeout = TIMEOUT if timeout is None else timeout
        kwargs.update({"timeout": timeout})

        func_name = sys._getframe().f_code.co_name
        res, check = api_request([client.unpin_snapshot_data, pin_id], **kwargs)
        check_result = ResponseChecker(res, func_name, check_task, check_items, check, pin_id=pin_id, **kwargs).run()
        return res, check_result