From a3cc6df3fd45fa2378a4980786b2408422ba5f53 Mon Sep 17 00:00:00 2001 From: PowderLi Date: Tue, 2 Jul 2024 20:45:22 +0800 Subject: [PATCH] fix: user-friendly 1. param filter is not requried 2. param limit is useless while outputFields = [count(*)] Signed-off-by: PowderLi --- .../distributed/proxy/httpserver/handler_v2.go | 8 +++++++- .../proxy/httpserver/handler_v2_test.go | 15 ++++++++++++++- .../distributed/proxy/httpserver/request_v2.go | 2 +- 3 files changed, 22 insertions(+), 3 deletions(-) diff --git a/internal/distributed/proxy/httpserver/handler_v2.go b/internal/distributed/proxy/httpserver/handler_v2.go index e8637a02b653f..9e5f6aa9cfac2 100644 --- a/internal/distributed/proxy/httpserver/handler_v2.go +++ b/internal/distributed/proxy/httpserver/handler_v2.go @@ -7,6 +7,7 @@ import ( "io" "net/http" "strconv" + "strings" "github.com/cockroachdb/errors" "github.com/gin-gonic/gin" @@ -563,6 +564,11 @@ func (h *HandlersV2) releaseCollection(ctx context.Context, c *gin.Context, anyR return resp, err } +// copy from internal/proxy/task_query.go +func matchCountRule(outputs []string) bool { + return len(outputs) == 1 && strings.ToLower(strings.TrimSpace(outputs[0])) == "count(*)" +} + func (h *HandlersV2) query(ctx context.Context, c *gin.Context, anyReq any, dbName string) (interface{}, error) { httpReq := anyReq.(*QueryReqV2) req := &milvuspb.QueryRequest{ @@ -578,7 +584,7 @@ func (h *HandlersV2) query(ctx context.Context, c *gin.Context, anyReq any, dbNa if httpReq.Offset > 0 { req.QueryParams = append(req.QueryParams, &commonpb.KeyValuePair{Key: ParamOffset, Value: strconv.FormatInt(int64(httpReq.Offset), 10)}) } - if httpReq.Limit > 0 { + if httpReq.Limit > 0 && !matchCountRule(httpReq.OutputFields) { req.QueryParams = append(req.QueryParams, &commonpb.KeyValuePair{Key: ParamLimit, Value: strconv.FormatInt(int64(httpReq.Limit), 10)}) } resp, err := wrapperProxy(ctx, c, req, h.checkAuth, false, func(reqCtx context.Context, req any) (interface{}, error) { diff --git a/internal/distributed/proxy/httpserver/handler_v2_test.go b/internal/distributed/proxy/httpserver/handler_v2_test.go index 5a57e536b963c..cd53733ddd88f 100644 --- a/internal/distributed/proxy/httpserver/handler_v2_test.go +++ b/internal/distributed/proxy/httpserver/handler_v2_test.go @@ -1159,7 +1159,16 @@ func TestDML(t *testing.T) { Status: &StatusSuccess, }, nil).Times(6) mp.EXPECT().DescribeCollection(mock.Anything, mock.Anything).Return(&milvuspb.DescribeCollectionResponse{Status: commonErrorStatus}, nil).Times(4) - mp.EXPECT().Query(mock.Anything, mock.Anything).Return(&milvuspb.QueryResults{Status: commonSuccessStatus, OutputFields: []string{}, FieldsData: []*schemapb.FieldData{}}, nil).Times(3) + mp.EXPECT().Query(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, req *milvuspb.QueryRequest) (*milvuspb.QueryResults, error) { + if matchCountRule(req.OutputFields) { + for _, pair := range req.QueryParams { + if pair.GetKey() == ParamLimit { + return nil, fmt.Errorf("mock error") + } + } + } + return &milvuspb.QueryResults{Status: commonSuccessStatus, OutputFields: []string{}, FieldsData: []*schemapb.FieldData{}}, nil + }).Times(4) mp.EXPECT().Insert(mock.Anything, mock.Anything).Return(&milvuspb.MutationResult{Status: commonSuccessStatus, InsertCnt: int64(0), IDs: &schemapb.IDs{IdField: &schemapb.IDs_IntId{IntId: &schemapb.LongArray{Data: []int64{}}}}}, nil).Once() mp.EXPECT().Insert(mock.Anything, mock.Anything).Return(&milvuspb.MutationResult{Status: commonSuccessStatus, InsertCnt: int64(0), IDs: &schemapb.IDs{IdField: &schemapb.IDs_StrId{StrId: &schemapb.StringArray{Data: []string{}}}}}, nil).Once() mp.EXPECT().Upsert(mock.Anything, mock.Anything).Return(&milvuspb.MutationResult{Status: commonSuccessStatus, UpsertCnt: int64(0), IDs: &schemapb.IDs{IdField: &schemapb.IDs_IntId{IntId: &schemapb.LongArray{Data: []int64{}}}}}, nil).Once() @@ -1181,6 +1190,10 @@ func TestDML(t *testing.T) { path: QueryAction, requestBody: []byte(`{"collectionName": "book", "filter": "book_id in [2, 4, 6, 8]"}`), }) + queryTestCases = append(queryTestCases, requestBodyTestCase{ + path: QueryAction, + requestBody: []byte(`{"collectionName": "book", "filter": "", "outputFields": ["count(*)"], "limit": 10}`), + }) queryTestCases = append(queryTestCases, requestBodyTestCase{ path: InsertAction, requestBody: []byte(`{"collectionName": "book", "data": [{"book_id": 0, "word_count": 0, "book_intro": [0.11825, 0.6]}]}`), diff --git a/internal/distributed/proxy/httpserver/request_v2.go b/internal/distributed/proxy/httpserver/request_v2.go index f5fe86a69e4d3..5f1ff0c08e0c4 100644 --- a/internal/distributed/proxy/httpserver/request_v2.go +++ b/internal/distributed/proxy/httpserver/request_v2.go @@ -104,7 +104,7 @@ type QueryReqV2 struct { CollectionName string `json:"collectionName" binding:"required"` PartitionNames []string `json:"partitionNames"` OutputFields []string `json:"outputFields"` - Filter string `json:"filter" binding:"required"` + Filter string `json:"filter"` Limit int32 `json:"limit"` Offset int32 `json:"offset"` }