grpc_args_parser.py 2.5 KB
Newer Older
P
peng.xu 已提交
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102
from milvus import Status
from functools import wraps


def error_status(func):
    @wraps(func)
    def inner(*args, **kwargs):
        try:
            results = func(*args, **kwargs)
        except Exception as e:
            return Status(code=Status.UNEXPECTED_ERROR, message=str(e)), None

        return Status(code=0, message="Success"), results

    return inner


class GrpcArgsParser(object):

    @classmethod
    @error_status
    def parse_proto_TableSchema(cls, param):
        _table_schema = {
            'status': param.status,
            'table_name': param.table_name,
            'dimension': param.dimension,
            'index_file_size': param.index_file_size,
            'metric_type': param.metric_type
        }

        return _table_schema

    @classmethod
    @error_status
    def parse_proto_TableName(cls, param):
        return param.table_name

    @classmethod
    @error_status
    def parse_proto_Index(cls, param):
        _index = {
            'index_type': param.index_type,
            'nlist': param.nlist
        }

        return _index

    @classmethod
    @error_status
    def parse_proto_IndexParam(cls, param):
        _table_name = param.table_name
        _status, _index = cls.parse_proto_Index(param.index)

        if not _status.OK():
            raise Exception("Argument parse error")

        return _table_name, _index

    @classmethod
    @error_status
    def parse_proto_Command(cls, param):
        _cmd = param.cmd

        return _cmd

    @classmethod
    @error_status
    def parse_proto_Range(cls, param):
        _start_value = param.start_value
        _end_value = param.end_value

        return _start_value, _end_value

    @classmethod
    @error_status
    def parse_proto_RowRecord(cls, param):
        return list(param.vector_data)

    @classmethod
    @error_status
    def parse_proto_SearchParam(cls, param):
        _table_name = param.table_name
        _topk = param.topk
        _nprobe = param.nprobe
        _status, _range = cls.parse_proto_Range(param.query_range_array)

        if not _status.OK():
            raise Exception("Argument parse error")

        _row_record = param.query_record_array

        return _table_name, _row_record, _range, _topk

    @classmethod
    @error_status
    def parse_proto_DeleteByRangeParam(cls, param):
        _table_name = param.table_name
        _range = param.range
        _start_value = _range.start_value
        _end_value = _range.end_value

        return _table_name, _start_value, _end_value