diff --git a/internal/proxy/rate_limit_interceptor.go b/internal/proxy/rate_limit_interceptor.go index f1d3cb0d75..b1184069ce 100644 --- a/internal/proxy/rate_limit_interceptor.go +++ b/internal/proxy/rate_limit_interceptor.go @@ -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: diff --git a/internal/proxy/rate_limit_interceptor_test.go b/internal/proxy/rate_limit_interceptor_test.go index d2149fc106..f6da7e43c1 100644 --- a/internal/proxy/rate_limit_interceptor_test.go +++ b/internal/proxy/rate_limit_interceptor_test.go @@ -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)