From d8e68cbb973d2b3e6810a01395e05bbd401e43ef Mon Sep 17 00:00:00 2001 From: SimFG Date: Fri, 12 Jul 2024 16:17:35 +0800 Subject: [PATCH] enhance: use the insert rate limit for the upsert request (#34615) - issue: #34614 Signed-off-by: SimFG --- internal/proxy/rate_limit_interceptor.go | 2 +- internal/proxy/rate_limit_interceptor_test.go | 2 +- pkg/common/common.go | 1 + 3 files changed, 3 insertions(+), 2 deletions(-) diff --git a/internal/proxy/rate_limit_interceptor.go b/internal/proxy/rate_limit_interceptor.go index 01030fb8f190b..0185237ea1b49 100644 --- a/internal/proxy/rate_limit_interceptor.go +++ b/internal/proxy/rate_limit_interceptor.go @@ -134,7 +134,7 @@ func getRequestInfo(ctx context.Context, req interface{}) (int64, map[int64][]in return dbID, collToPartIDs, internalpb.RateType_DMLInsert, proto.Size(r), err case *milvuspb.UpsertRequest: dbID, collToPartIDs, err := getCollectionAndPartitionID(ctx, req.(reqPartName)) - return dbID, collToPartIDs, internalpb.RateType_DMLUpsert, proto.Size(r), err + return dbID, collToPartIDs, internalpb.RateType_DMLInsert, proto.Size(r), err case *milvuspb.DeleteRequest: dbID, collToPartIDs, err := getCollectionAndPartitionID(ctx, req.(reqPartName)) return dbID, collToPartIDs, internalpb.RateType_DMLDelete, proto.Size(r), err diff --git a/internal/proxy/rate_limit_interceptor_test.go b/internal/proxy/rate_limit_interceptor_test.go index 3f9fbac8435a2..9004bc8d56858 100644 --- a/internal/proxy/rate_limit_interceptor_test.go +++ b/internal/proxy/rate_limit_interceptor_test.go @@ -96,7 +96,7 @@ func TestRateLimitInterceptor(t *testing.T) { PartitionName: "p1", DbName: "db1", }), size) - assert.Equal(t, internalpb.RateType_DMLUpsert, rt) + assert.Equal(t, internalpb.RateType_DMLInsert, rt) assert.Equal(t, database, int64(100)) assert.True(t, len(col2part) == 1) assert.Equal(t, int64(10), col2part[1][0]) diff --git a/pkg/common/common.go b/pkg/common/common.go index 006c77f1e12b1..bf51e769294ae 100644 --- a/pkg/common/common.go +++ b/pkg/common/common.go @@ -23,6 +23,7 @@ import ( "strings" "github.com/cockroachdb/errors" + "github.com/milvus-io/milvus-proto/go-api/v2/commonpb" "github.com/milvus-io/milvus-proto/go-api/v2/schemapb" )