From 3a7154b796ea0739a46605c206a80fee162ccbec Mon Sep 17 00:00:00 2001 From: SimFG Date: Mon, 29 Apr 2024 10:19:31 +0800 Subject: [PATCH] enhance: [2.3] add the skip auto id and partition key check config (#32671) /kind improvement issue: #32591 pr: #32592 Signed-off-by: SimFG --- internal/proxy/util.go | 8 ++++++-- pkg/util/paramtable/component_param.go | 19 +++++++++++++++++++ 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 8d1f585b58..5bbdca8df4 100644 --- a/internal/proxy/util.go +++ b/internal/proxy/util.go @@ -1207,7 +1207,11 @@ 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)) @@ -1256,7 +1260,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 c0324381e8..b105f8107c 100644 --- a/pkg/util/paramtable/component_param.go +++ b/pkg/util/paramtable/component_param.go @@ -979,6 +979,9 @@ type proxyConfig struct { RetryTimesOnHealthCheck ParamItem `refreshable:"true"` PartitionNameRegexp ParamItem `refreshable:"true"` + SkipAutoIDCheck ParamItem `refreshable:"true"` + SkipPartitionKeyCheck ParamItem `refreshable:"true"` + AccessLog AccessLogConfig // connection manager @@ -1282,6 +1285,22 @@ please adjust in embedded Milvus: false`, } p.PartitionNameRegexp.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 0d6a66fca5..937ea68f5a 100644 --- a/pkg/util/paramtable/component_param_test.go +++ b/pkg/util/paramtable/component_param_test.go @@ -173,6 +173,14 @@ func TestComponentParam(t *testing.T) { params.Save("proxy.gracefulStopTimeout", "100") assert.Equal(t, 100*time.Second, Params.GracefulStopTimeout.GetAsDuration(time.Second)) + + 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) {