milvus-io
/
milvus
mirror of https://github.com/milvus-io/milvus.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

enhance: [2.4] add the cost response for the rest api (#32841) 

issue: #30436
pr: #32620

Signed-off-by: SimFG <bang.fu@zilliz.com>
pull/32846/head
SimFG 2024-05-08 14:19:29 +08:00 committed by GitHub
parent 3818c166c5
commit 5023f1a446
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
5 changed files with 99 additions and 13 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
internal/distributed/proxy/httpserver/constant.go
View File

@ -87,6 +87,7 @@ const (
HTTPReturnCode = "code"
HTTPReturnMessage = "message"
HTTPReturnData = "data"
HTTPReturnCost = "cost"
HTTPReturnLoadState = "loadState"
HTTPReturnLoadProgress = "loadProgress"

64
internal/distributed/proxy/httpserver/handler_v2.go
View File

@ -586,7 +586,11 @@ func (h *HandlersV2) query(ctx context.Context, c *gin.Context, anyReq any, dbNa
HTTPReturnMessage: merr.ErrInvalidSearchResult.Error() + ", error: " + err.Error(),
})
} else {
c.JSON(http.StatusOK, gin.H{HTTPReturnCode: http.StatusOK, HTTPReturnData: outputData})
c.JSON(http.StatusOK, gin.H{
HTTPReturnCode: http.StatusOK,
HTTPReturnData: outputData,
HTTPReturnCost: proxy.GetCostValue(queryResp.GetStatus()),
})
}
}
return resp, err
@ -629,7 +633,11 @@ func (h *HandlersV2) get(ctx context.Context, c *gin.Context, anyReq any, dbName
HTTPReturnMessage: merr.ErrInvalidSearchResult.Error() + ", error: " + err.Error(),
})
} else {
c.JSON(http.StatusOK, gin.H{HTTPReturnCode: http.StatusOK, HTTPReturnData: outputData})
c.JSON(http.StatusOK, gin.H{
HTTPReturnCode: http.StatusOK,
HTTPReturnData: outputData,
HTTPReturnCost: proxy.GetCostValue(queryResp.GetStatus()),
})
}
}
return resp, err
@ -663,7 +671,9 @@ func (h *HandlersV2) delete(ctx context.Context, c *gin.Context, anyReq any, dbN
return h.proxy.Delete(reqCtx, req.(*milvuspb.DeleteRequest))
})
if err == nil {
c.JSON(http.StatusOK, wrapperReturnDefault())
c.JSON(http.StatusOK, wrapperReturnDefaultWithCost(
proxy.GetCostValue(resp.(*milvuspb.MutationResult).GetStatus()),
))
}
return resp, err
}
@ -705,16 +715,29 @@ func (h *HandlersV2) insert(ctx context.Context, c *gin.Context, anyReq any, dbN
})
if err == nil {
insertResp := resp.(*milvuspb.MutationResult)
cost := proxy.GetCostValue(insertResp.GetStatus())
switch insertResp.IDs.GetIdField().(type) {
case *schemapb.IDs_IntId:
allowJS, _ := strconv.ParseBool(c.Request.Header.Get(HTTPHeaderAllowInt64))
if allowJS {
c.JSON(http.StatusOK, gin.H{HTTPReturnCode: http.StatusOK, HTTPReturnData: gin.H{"insertCount": insertResp.InsertCnt, "insertIds": insertResp.IDs.IdField.(*schemapb.IDs_IntId).IntId.Data}})
c.JSON(http.StatusOK, gin.H{
HTTPReturnCode: http.StatusOK,
HTTPReturnData: gin.H{"insertCount": insertResp.InsertCnt, "insertIds": insertResp.IDs.IdField.(*schemapb.IDs_IntId).IntId.Data},
HTTPReturnCost: cost,
})
} else {
c.JSON(http.StatusOK, gin.H{HTTPReturnCode: http.StatusOK, HTTPReturnData: gin.H{"insertCount": insertResp.InsertCnt, "insertIds": formatInt64(insertResp.IDs.IdField.(*schemapb.IDs_IntId).IntId.Data)}})
c.JSON(http.StatusOK, gin.H{
HTTPReturnCode: http.StatusOK,
HTTPReturnData: gin.H{"insertCount": insertResp.InsertCnt, "insertIds": formatInt64(insertResp.IDs.IdField.(*schemapb.IDs_IntId).IntId.Data)},
HTTPReturnCost: cost,
})
}
case *schemapb.IDs_StrId:
c.JSON(http.StatusOK, gin.H{HTTPReturnCode: http.StatusOK, HTTPReturnData: gin.H{"insertCount": insertResp.InsertCnt, "insertIds": insertResp.IDs.IdField.(*schemapb.IDs_StrId).StrId.Data}})
c.JSON(http.StatusOK, gin.H{
HTTPReturnCode: http.StatusOK,
HTTPReturnData: gin.H{"insertCount": insertResp.InsertCnt, "insertIds": insertResp.IDs.IdField.(*schemapb.IDs_StrId).StrId.Data},
HTTPReturnCost: cost,
})
default:
c.JSON(http.StatusOK, gin.H{
HTTPReturnCode: merr.Code(merr.ErrCheckPrimaryKey),
@ -767,16 +790,29 @@ func (h *HandlersV2) upsert(ctx context.Context, c *gin.Context, anyReq any, dbN
})
if err == nil {
upsertResp := resp.(*milvuspb.MutationResult)
cost := proxy.GetCostValue(upsertResp.GetStatus())
switch upsertResp.IDs.GetIdField().(type) {
case *schemapb.IDs_IntId:
allowJS, _ := strconv.ParseBool(c.Request.Header.Get(HTTPHeaderAllowInt64))
if allowJS {
c.JSON(http.StatusOK, gin.H{HTTPReturnCode: http.StatusOK, HTTPReturnData: gin.H{"upsertCount": upsertResp.UpsertCnt, "upsertIds": upsertResp.IDs.IdField.(*schemapb.IDs_IntId).IntId.Data}})
c.JSON(http.StatusOK, gin.H{
HTTPReturnCode: http.StatusOK,
HTTPReturnData: gin.H{"upsertCount": upsertResp.UpsertCnt, "upsertIds": upsertResp.IDs.IdField.(*schemapb.IDs_IntId).IntId.Data},
HTTPReturnCost: cost,
})
} else {
c.JSON(http.StatusOK, gin.H{HTTPReturnCode: http.StatusOK, HTTPReturnData: gin.H{"upsertCount": upsertResp.UpsertCnt, "upsertIds": formatInt64(upsertResp.IDs.IdField.(*schemapb.IDs_IntId).IntId.Data)}})
c.JSON(http.StatusOK, gin.H{
HTTPReturnCode: http.StatusOK,
HTTPReturnData: gin.H{"upsertCount": upsertResp.UpsertCnt, "upsertIds": formatInt64(upsertResp.IDs.IdField.(*schemapb.IDs_IntId).IntId.Data)},
HTTPReturnCost: cost,
})
}
case *schemapb.IDs_StrId:
c.JSON(http.StatusOK, gin.H{HTTPReturnCode: http.StatusOK, HTTPReturnData: gin.H{"upsertCount": upsertResp.UpsertCnt, "upsertIds": upsertResp.IDs.IdField.(*schemapb.IDs_StrId).StrId.Data}})
c.JSON(http.StatusOK, gin.H{
HTTPReturnCode: http.StatusOK,
HTTPReturnData: gin.H{"upsertCount": upsertResp.UpsertCnt, "upsertIds": upsertResp.IDs.IdField.(*schemapb.IDs_StrId).StrId.Data},
HTTPReturnCost: cost,
})
default:
c.JSON(http.StatusOK, gin.H{
HTTPReturnCode: merr.Code(merr.ErrCheckPrimaryKey),
@ -894,8 +930,9 @@ func (h *HandlersV2) search(ctx context.Context, c *gin.Context, anyReq any, dbN
})
if err == nil {
searchResp := resp.(*milvuspb.SearchResults)
cost := proxy.GetCostValue(searchResp.GetStatus())
if searchResp.Results.TopK == int64(0) {
c.JSON(http.StatusOK, gin.H{HTTPReturnCode: http.StatusOK, HTTPReturnData: []interface{}{}})
c.JSON(http.StatusOK, gin.H{HTTPReturnCode: http.StatusOK, HTTPReturnData: []interface{}{}, HTTPReturnCost: cost})
} else {
allowJS, _ := strconv.ParseBool(c.Request.Header.Get(HTTPHeaderAllowInt64))
outputData, err := buildQueryResp(searchResp.Results.TopK, searchResp.Results.OutputFields, searchResp.Results.FieldsData, searchResp.Results.Ids, searchResp.Results.Scores, allowJS)
@ -906,7 +943,7 @@ func (h *HandlersV2) search(ctx context.Context, c *gin.Context, anyReq any, dbN
HTTPReturnMessage: merr.ErrInvalidSearchResult.Error() + ", error: " + err.Error(),
})
} else {
c.JSON(http.StatusOK, gin.H{HTTPReturnCode: http.StatusOK, HTTPReturnData: outputData})
c.JSON(http.StatusOK, gin.H{HTTPReturnCode: http.StatusOK, HTTPReturnData: outputData, HTTPReturnCost: cost})
}
}
}
@ -971,8 +1008,9 @@ func (h *HandlersV2) advancedSearch(ctx context.Context, c *gin.Context, anyReq
})
if err == nil {
searchResp := resp.(*milvuspb.SearchResults)
cost := proxy.GetCostValue(searchResp.GetStatus())
if searchResp.Results.TopK == int64(0) {
c.JSON(http.StatusOK, gin.H{HTTPReturnCode: http.StatusOK, HTTPReturnData: []interface{}{}})
c.JSON(http.StatusOK, gin.H{HTTPReturnCode: http.StatusOK, HTTPReturnData: []interface{}{}, HTTPReturnCost: cost})
} else {
allowJS, _ := strconv.ParseBool(c.Request.Header.Get(HTTPHeaderAllowInt64))
outputData, err := buildQueryResp(0, searchResp.Results.OutputFields, searchResp.Results.FieldsData, searchResp.Results.Ids, searchResp.Results.Scores, allowJS)
@ -983,7 +1021,7 @@ func (h *HandlersV2) advancedSearch(ctx context.Context, c *gin.Context, anyReq
HTTPReturnMessage: merr.ErrInvalidSearchResult.Error() + ", error: " + err.Error(),
})
} else {
c.JSON(http.StatusOK, gin.H{HTTPReturnCode: http.StatusOK, HTTPReturnData: outputData})
c.JSON(http.StatusOK, gin.H{HTTPReturnCode: http.StatusOK, HTTPReturnData: outputData, HTTPReturnCost: cost})
}
}
}

4
internal/distributed/proxy/httpserver/request_v2.go
View File

@ -379,3 +379,7 @@ func wrapperReturnRowCount(pairs []*commonpb.KeyValuePair) gin.H {
func wrapperReturnDefault() gin.H {
return gin.H{HTTPReturnCode: http.StatusOK, HTTPReturnData: gin.H{}}
}
func wrapperReturnDefaultWithCost(cost int) gin.H {
return gin.H{HTTPReturnCode: http.StatusOK, HTTPReturnData: gin.H{}, HTTPReturnCost: cost}
}

11
internal/proxy/util.go
View File

@ -1642,3 +1642,14 @@ func SetReportValue(status *commonpb.Status, value int) {
}
status.ExtraInfo["report_value"] = strconv.Itoa(value)
}
func GetCostValue(status *commonpb.Status) int {
if status == nil || status.ExtraInfo == nil {
return 0
}
value, err := strconv.Atoi(status.ExtraInfo["report_value"])
if err != nil {
return 0
}
return value
}

32
internal/proxy/util_test.go
View File

@ -2209,3 +2209,35 @@ func TestAppendUserInfoForRPC(t *testing.T) {
expectAuth := crypto.Base64Encode("root:root")
assert.Equal(t, expectAuth, authorization[0])
}
func TestGetCostValue(t *testing.T) {
t.Run("empty status", func(t *testing.T) {
{
cost := GetCostValue(&commonpb.Status{})
assert.Equal(t, 0, cost)
}
{
cost := GetCostValue(nil)
assert.Equal(t, 0, cost)
}
})
t.Run("wrong cost value style", func(t *testing.T) {
cost := GetCostValue(&commonpb.Status{
ExtraInfo: map[string]string{
"report_value": "abc",
},
})
assert.Equal(t, 0, cost)
})
t.Run("success", func(t *testing.T) {
cost := GetCostValue(&commonpb.Status{
ExtraInfo: map[string]string{
"report_value": "100",
},
})
assert.Equal(t, 100, cost)
})
}