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

Intercept import request in Proxy interceptor (#19856) 

Signed-off-by: bigsheeper <yihao.dai@zilliz.com>

Signed-off-by: bigsheeper <yihao.dai@zilliz.com>
pull/19871/head
bigsheeper 2022-10-18 16:53:26 +08:00 committed by GitHub
parent 2785f98de6
commit 15f7592739
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 7 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
internal/proxy/rate_limit_interceptor.go
View File

@ -60,7 +60,8 @@ func getRequestInfo(req interface{}) (internalpb.RateType, int, error) {
return internalpb.RateType_DMLInsert, proto.Size(r), nil
case *milvuspb.DeleteRequest:
return internalpb.RateType_DMLDelete, proto.Size(r), nil
// TODO: add bulkLoad
case *milvuspb.ImportRequest:
return internalpb.RateType_DMLBulkLoad, proto.Size(r), nil
case *milvuspb.SearchRequest:
return internalpb.RateType_DQLSearch, int(r.GetNq()), nil
case *milvuspb.QueryRequest:

5
internal/proxy/rate_limit_interceptor_test.go
View File

@ -50,6 +50,11 @@ func TestRateLimitInterceptor(t *testing.T) {
assert.Equal(t, proto.Size(&milvuspb.DeleteRequest{}), size)
assert.Equal(t, internalpb.RateType_DMLDelete, rt)
rt, size, err = getRequestInfo(&milvuspb.ImportRequest{})
assert.NoError(t, err)
assert.Equal(t, proto.Size(&milvuspb.ImportRequest{}), size)
assert.Equal(t, internalpb.RateType_DMLBulkLoad, rt)
rt, size, err = getRequestInfo(&milvuspb.SearchRequest{})
assert.NoError(t, err)
assert.Equal(t, proto.Size(&milvuspb.SearchRequest{}), size)