enhance: use the insert rate limit for the upsert request (#34615)

- issue: #34614 Signed-off-by: SimFG <bang.fu@zilliz.com>
milvus-io · Jul 12, 2024 · d8e68cb · d8e68cb
1 parent 531092c
commit d8e68cb
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 2 deletions.
diff --git 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
@@ -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
@@ -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"
 )