enhance: [Restful] Add consistency level for query/get API (#41825)

Related to #41805

Signed-off-by: Congqi Xia <congqi.xia@zilliz.com>
This commit is contained in:
congqixia 2025-05-14 10:48:21 +08:00 committed by GitHub
parent 21d6d1669e
commit fb612c765c
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 61 additions and 30 deletions

View File

@ -788,7 +788,16 @@ func (h *HandlersV2) query(ctx context.Context, c *gin.Context, anyReq any, dbNa
OutputFields: httpReq.OutputFields,
PartitionNames: httpReq.PartitionNames,
QueryParams: []*commonpb.KeyValuePair{},
UseDefaultConsistency: true,
}
var err error
req.ConsistencyLevel, req.UseDefaultConsistency, err = convertConsistencyLevel(httpReq.ConsistencyLevel)
if err != nil {
log.Ctx(ctx).Warn("high level restful api, query with consistency_level invalid", zap.Error(err))
HTTPAbortReturn(c, http.StatusOK, gin.H{
HTTPReturnCode: merr.Code(err),
HTTPReturnMessage: "consistencyLevel can only be [Strong, Session, Bounded, Eventually, Customized], default: Bounded, err:" + err.Error(),
})
return nil, err
}
req.ExprTemplateValues = generateExpressionTemplate(httpReq.ExprParams)
c.Set(ContextRequest, req)
@ -843,7 +852,15 @@ func (h *HandlersV2) get(ctx context.Context, c *gin.Context, anyReq any, dbName
OutputFields: httpReq.OutputFields,
PartitionNames: httpReq.PartitionNames,
Expr: filter,
UseDefaultConsistency: true,
}
req.ConsistencyLevel, req.UseDefaultConsistency, err = convertConsistencyLevel(httpReq.ConsistencyLevel)
if err != nil {
log.Ctx(ctx).Warn("high level restful api, query with consistency_level invalid", zap.Error(err))
HTTPAbortReturn(c, http.StatusOK, gin.H{
HTTPReturnCode: merr.Code(err),
HTTPReturnMessage: "consistencyLevel can only be [Strong, Session, Bounded, Eventually, Customized], default: Bounded, err:" + err.Error(),
})
return nil, err
}
c.Set(ContextRequest, req)
resp, err := wrapperProxyWithLimit(ctx, c, req, h.checkAuth, false, "/milvus.proto.milvus.MilvusService/Query", true, h.proxy, func(reqCtx context.Context, req any) (interface{}, error) {

View File

@ -2135,7 +2135,7 @@ func TestDML(t *testing.T) {
Schema: generateCollectionSchema(schemapb.DataType_Int64, false, true),
ShardsNum: ShardNumDefault,
Status: &StatusSuccess,
}, nil).Times(6)
}, nil).Times(7)
mp.EXPECT().DescribeCollection(mock.Anything, mock.Anything).Return(&milvuspb.DescribeCollectionResponse{Status: commonErrorStatus}, nil).Times(4)
mp.EXPECT().Query(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, req *milvuspb.QueryRequest) (*milvuspb.QueryResults, error) {
if matchCountRule(req.OutputFields) {
@ -2179,6 +2179,18 @@ func TestDML(t *testing.T) {
path: QueryAction,
requestBody: []byte(`{"collectionName": "book", "filter": "", "outputFields": ["count(*)"], "limit": 10}`),
})
queryTestCases = append(queryTestCases, requestBodyTestCase{
path: QueryAction,
requestBody: []byte(`{"collectionName": "book", "filter": "", "outputFields": ["book_id", "word_count", "book_intro"], "limit": 10, "consistencyLevel": "AAA"}`),
errMsg: "consistencyLevel can only be [Strong, Session, Bounded, Eventually, Customized], default: Bounded",
errCode: 1100, // ErrParameterInvalid
})
queryTestCases = append(queryTestCases, requestBodyTestCase{
path: GetAction,
requestBody: []byte(`{"collectionName": "book", "id": [2, 4, 6, 8], "outputFields": ["book_id", "word_count", "book_intro"], "consistencyLevel": "AAA"}`),
errMsg: "consistencyLevel can only be [Strong, Session, Bounded, Eventually, Customized], default: Bounded",
errCode: 1100, // ErrParameterInvalid
})
queryTestCases = append(queryTestCases, requestBodyTestCase{
path: InsertAction,
requestBody: []byte(`{"collectionName": "book", "data": [{"book_id": 0, "word_count": 0, "book_intro": [0.11825, 0.6]}]}`),
@ -2211,7 +2223,7 @@ func TestDML(t *testing.T) {
})
queryTestCases = append(queryTestCases, requestBodyTestCase{
path: GetAction,
requestBody: []byte(`{"collectionName": "book", "id" : [2, 4, 6, 8, 0], "outputFields": ["book_id", "word_count", "book_intro"]}`),
requestBody: []byte(`{"collectionName": "book", "id": [2, 4, 6, 8, "0"], "outputFields": ["book_id", "word_count", "book_intro"]}`),
errMsg: "",
errCode: 65535,
})

View File

@ -217,6 +217,7 @@ type QueryReqV2 struct {
Limit int32 `json:"limit"`
Offset int32 `json:"offset"`
ExprParams map[string]interface{} `json:"exprParams"`
ConsistencyLevel string `json:"consistencyLevel"`
}
func (req *QueryReqV2) GetDbName() string { return req.DbName }
@ -228,6 +229,7 @@ type CollectionIDReq struct {
PartitionNames []string `json:"partitionNames"`
OutputFields []string `json:"outputFields"`
ID interface{} `json:"id" binding:"required"`
ConsistencyLevel string `json:"consistencyLevel"`
}
func (req *CollectionIDReq) GetDbName() string { return req.DbName }