From 09cd56d44f0739ab1b8d440cd1b84e30798e7265 Mon Sep 17 00:00:00 2001 From: SimFG Date: Mon, 29 Apr 2024 10:29:26 +0800 Subject: [PATCH] enhance: add the skip auto id and partition key check config (#32592) /kind improvement issue: #32591 Signed-off-by: SimFG --- internal/proxy/util.go | 9 +++++++-- pkg/util/paramtable/component_param.go | 18 ++++++++++++++++++ pkg/util/paramtable/component_param_test.go | 8 ++++++++ 3 files changed, 33 insertions(+), 2 deletions(-) diff --git a/internal/proxy/util.go b/internal/proxy/util.go index 849791da64..586f759c35 100644 --- a/internal/proxy/util.go +++ b/internal/proxy/util.go @@ -1190,7 +1190,12 @@ func checkPrimaryFieldData(schema *schemapb.CollectionSchema, result *milvuspb.M var primaryFieldData *schemapb.FieldData if inInsert { // when checkPrimaryFieldData in insert - if !primaryFieldSchema.AutoID { + + skipAutoIDCheck := Params.ProxyCfg.SkipAutoIDCheck.GetAsBool() && + primaryFieldSchema.AutoID && + typeutil.IsPrimaryFieldDataExist(insertMsg.GetFieldsData(), primaryFieldSchema) + + if !primaryFieldSchema.AutoID || skipAutoIDCheck { primaryFieldData, err = typeutil.GetPrimaryFieldData(insertMsg.GetFieldsData(), primaryFieldSchema) if err != nil { log.Info("get primary field data failed", zap.String("collectionName", insertMsg.CollectionName), zap.Error(err)) @@ -1239,7 +1244,7 @@ func checkPrimaryFieldData(schema *schemapb.CollectionSchema, result *milvuspb.M } func getPartitionKeyFieldData(fieldSchema *schemapb.FieldSchema, insertMsg *msgstream.InsertMsg) (*schemapb.FieldData, error) { - if len(insertMsg.GetPartitionName()) > 0 { + if len(insertMsg.GetPartitionName()) > 0 && !Params.ProxyCfg.SkipPartitionKeyCheck.GetAsBool() { return nil, errors.New("not support manually specifying the partition names if partition key mode is used") } diff --git a/pkg/util/paramtable/component_param.go b/pkg/util/paramtable/component_param.go index 0a35f3508f..a900fb12b7 100644 --- a/pkg/util/paramtable/component_param.go +++ b/pkg/util/paramtable/component_param.go @@ -1012,6 +1012,8 @@ type proxyConfig struct { RetryTimesOnHealthCheck ParamItem `refreshable:"true"` PartitionNameRegexp ParamItem `refreshable:"true"` MustUsePartitionKey ParamItem `refreshable:"true"` + SkipAutoIDCheck ParamItem `refreshable:"true"` + SkipPartitionKeyCheck ParamItem `refreshable:"true"` AccessLog AccessLogConfig @@ -1348,6 +1350,22 @@ please adjust in embedded Milvus: false`, } p.MustUsePartitionKey.Init(base.mgr) + p.SkipAutoIDCheck = ParamItem{ + Key: "proxy.skipAutoIDCheck", + Version: "2.4.1", + DefaultValue: "false", + Doc: "switch for whether proxy shall skip auto id check when inserting data", + } + p.SkipAutoIDCheck.Init(base.mgr) + + p.SkipPartitionKeyCheck = ParamItem{ + Key: "proxy.skipPartitionKeyCheck", + Version: "2.4.1", + DefaultValue: "false", + Doc: "switch for whether proxy shall skip partition key check when inserting data", + } + p.SkipPartitionKeyCheck.Init(base.mgr) + p.GracefulStopTimeout = ParamItem{ Key: "proxy.gracefulStopTimeout", Version: "2.3.7", diff --git a/pkg/util/paramtable/component_param_test.go b/pkg/util/paramtable/component_param_test.go index 5627e6696c..f62ded1997 100644 --- a/pkg/util/paramtable/component_param_test.go +++ b/pkg/util/paramtable/component_param_test.go @@ -174,6 +174,14 @@ func TestComponentParam(t *testing.T) { assert.False(t, Params.MustUsePartitionKey.GetAsBool()) params.Save("proxy.mustUsePartitionKey", "true") assert.True(t, Params.MustUsePartitionKey.GetAsBool()) + + assert.False(t, Params.SkipAutoIDCheck.GetAsBool()) + params.Save("proxy.skipAutoIDCheck", "true") + assert.True(t, Params.SkipAutoIDCheck.GetAsBool()) + + assert.False(t, Params.SkipPartitionKeyCheck.GetAsBool()) + params.Save("proxy.skipPartitionKeyCheck", "true") + assert.True(t, Params.SkipPartitionKeyCheck.GetAsBool()) }) // t.Run("test proxyConfig panic", func(t *testing.T) {