From a8caa031e94702cbb41fdab47e295bc28bb2feea Mon Sep 17 00:00:00 2001 From: SimFG Date: Tue, 26 Nov 2024 22:44:31 +0800 Subject: [PATCH] enhance: [2.4] support to config the default db properties Signed-off-by: SimFG --- internal/metastore/model/database.go | 4 +-- internal/metastore/model/database_test.go | 2 +- .../rootcoord/create_collection_task_test.go | 4 +-- internal/rootcoord/create_db_task_test.go | 4 +-- internal/rootcoord/describe_db_task_test.go | 2 +- internal/rootcoord/list_db_task_test.go | 4 +-- internal/rootcoord/meta_table.go | 8 +++++- internal/rootcoord/meta_table_test.go | 26 +++++++++---------- pkg/util/funcutil/func.go | 9 +++++++ pkg/util/funcutil/func_test.go | 21 +++++++++++++++ pkg/util/paramtable/component_param.go | 10 +++++++ pkg/util/paramtable/component_param_test.go | 4 +++ 12 files changed, 74 insertions(+), 24 deletions(-) diff --git a/internal/metastore/model/database.go b/internal/metastore/model/database.go index 52fe49fb82243..95166a65b5de6 100644 --- a/internal/metastore/model/database.go +++ b/internal/metastore/model/database.go @@ -31,8 +31,8 @@ func NewDatabase(id int64, name string, state pb.DatabaseState, properties []*co } } -func NewDefaultDatabase() *Database { - return NewDatabase(util.DefaultDBID, util.DefaultDBName, pb.DatabaseState_DatabaseCreated, nil) +func NewDefaultDatabase(prop []*commonpb.KeyValuePair) *Database { + return NewDatabase(util.DefaultDBID, util.DefaultDBName, pb.DatabaseState_DatabaseCreated, prop) } func (c *Database) Available() bool { diff --git a/internal/metastore/model/database_test.go b/internal/metastore/model/database_test.go index 8effb2217acd2..622674d285171 100644 --- a/internal/metastore/model/database_test.go +++ b/internal/metastore/model/database_test.go @@ -59,5 +59,5 @@ func TestDatabaseCloneAndEqual(t *testing.T) { func TestDatabaseAvailable(t *testing.T) { assert.True(t, dbModel.Available()) - assert.True(t, NewDefaultDatabase().Available()) + assert.True(t, NewDefaultDatabase(nil).Available()) } diff --git a/internal/rootcoord/create_collection_task_test.go b/internal/rootcoord/create_collection_task_test.go index 6cc64bafd96ac..d923dd93ddc3d 100644 --- a/internal/rootcoord/create_collection_task_test.go +++ b/internal/rootcoord/create_collection_task_test.go @@ -636,7 +636,7 @@ func Test_createCollectionTask_Prepare(t *testing.T) { mock.Anything, mock.Anything, mock.Anything, - ).Return(model.NewDefaultDatabase(), nil) + ).Return(model.NewDefaultDatabase(nil), nil) meta.On("ListAllAvailCollections", mock.Anything, ).Return(map[int64][]int64{ @@ -1050,7 +1050,7 @@ func Test_createCollectionTask_PartitionKey(t *testing.T) { mock.Anything, mock.Anything, mock.Anything, - ).Return(model.NewDefaultDatabase(), nil) + ).Return(model.NewDefaultDatabase(nil), nil) meta.On("ListAllAvailCollections", mock.Anything, ).Return(map[int64][]int64{ diff --git a/internal/rootcoord/create_db_task_test.go b/internal/rootcoord/create_db_task_test.go index 09fb9c5a2397c..5f83139150d6e 100644 --- a/internal/rootcoord/create_db_task_test.go +++ b/internal/rootcoord/create_db_task_test.go @@ -54,7 +54,7 @@ func Test_CreateDBTask_Prepare(t *testing.T) { cfgMaxDatabaseNum := Params.RootCoordCfg.MaxDatabaseNum.GetAsInt() dbs := make([]*model.Database, 0, cfgMaxDatabaseNum) for i := 0; i < cfgMaxDatabaseNum; i++ { - dbs = append(dbs, model.NewDefaultDatabase()) + dbs = append(dbs, model.NewDefaultDatabase(nil)) } meta.On("ListDatabases", mock.Anything, @@ -81,7 +81,7 @@ func Test_CreateDBTask_Prepare(t *testing.T) { meta.On("ListDatabases", mock.Anything, mock.Anything). - Return([]*model.Database{model.NewDefaultDatabase()}, nil) + Return([]*model.Database{model.NewDefaultDatabase(nil)}, nil) core := newTestCore(withMeta(meta), withValidIDAllocator()) paramtable.Get().Save(Params.RootCoordCfg.MaxDatabaseNum.Key, strconv.Itoa(10)) diff --git a/internal/rootcoord/describe_db_task_test.go b/internal/rootcoord/describe_db_task_test.go index 9d86708d92f42..2d9b8a4f5ef32 100644 --- a/internal/rootcoord/describe_db_task_test.go +++ b/internal/rootcoord/describe_db_task_test.go @@ -48,7 +48,7 @@ func Test_describeDatabaseTask_Execute(t *testing.T) { t.Run("describe with empty database name", func(t *testing.T) { meta := mockrootcoord.NewIMetaTable(t) meta.EXPECT().GetDatabaseByName(mock.Anything, mock.Anything, mock.Anything). - Return(model.NewDefaultDatabase(), nil) + Return(model.NewDefaultDatabase(nil), nil) core := newTestCore(withMeta(meta)) task := &describeDBTask{ diff --git a/internal/rootcoord/list_db_task_test.go b/internal/rootcoord/list_db_task_test.go index 29d4feb3a62c5..db5ac702535fd 100644 --- a/internal/rootcoord/list_db_task_test.go +++ b/internal/rootcoord/list_db_task_test.go @@ -58,7 +58,7 @@ func Test_ListDBTask(t *testing.T) { }) t.Run("ok", func(t *testing.T) { - ret := []*model.Database{model.NewDefaultDatabase()} + ret := []*model.Database{model.NewDefaultDatabase(nil)} meta := mockrootcoord.NewIMetaTable(t) meta.On("ListDatabases", mock.Anything, @@ -89,7 +89,7 @@ func Test_ListDBTask(t *testing.T) { t.Run("list db with auth", func(t *testing.T) { Params.Save(Params.CommonCfg.AuthorizationEnabled.Key, "true") defer Params.Reset(Params.CommonCfg.AuthorizationEnabled.Key) - ret := []*model.Database{model.NewDefaultDatabase()} + ret := []*model.Database{model.NewDefaultDatabase(nil)} meta := mockrootcoord.NewIMetaTable(t) core := newTestCore(withMeta(meta)) diff --git a/internal/rootcoord/meta_table.go b/internal/rootcoord/meta_table.go index fd22fe7b0f64c..ba48c077c900f 100644 --- a/internal/rootcoord/meta_table.go +++ b/internal/rootcoord/meta_table.go @@ -258,7 +258,13 @@ func (mt *MetaTable) createDefaultDb() error { return err } - return mt.createDatabasePrivate(mt.ctx, model.NewDefaultDatabase(), ts) + s := Params.RootCoordCfg.DefaultDBProperties.GetValue() + defaultProperties, err := funcutil.String2KeyValuePair(s) + if err != nil { + return err + } + + return mt.createDatabasePrivate(mt.ctx, model.NewDefaultDatabase(defaultProperties), ts) } func (mt *MetaTable) CreateDatabase(ctx context.Context, db *model.Database, ts typeutil.Timestamp) error { diff --git a/internal/rootcoord/meta_table_test.go b/internal/rootcoord/meta_table_test.go index 3bf93aeb31cdb..67a3e7d8f9d68 100644 --- a/internal/rootcoord/meta_table_test.go +++ b/internal/rootcoord/meta_table_test.go @@ -460,7 +460,7 @@ func TestMetaTable_getCollectionByIDInternal(t *testing.T) { meta := &MetaTable{ catalog: catalog, dbName2Meta: map[string]*model.Database{ - util.DefaultDBName: model.NewDefaultDatabase(), + util.DefaultDBName: model.NewDefaultDatabase(nil), }, collID2Meta: map[typeutil.UniqueID]*model.Collection{}, } @@ -480,7 +480,7 @@ func TestMetaTable_getCollectionByIDInternal(t *testing.T) { meta := &MetaTable{ catalog: catalog, dbName2Meta: map[string]*model.Database{ - util.DefaultDBName: model.NewDefaultDatabase(), + util.DefaultDBName: model.NewDefaultDatabase(nil), }, collID2Meta: map[typeutil.UniqueID]*model.Collection{}, } @@ -618,7 +618,7 @@ func TestMetaTable_GetCollectionByName(t *testing.T) { ).Return(nil, errors.New("error mock GetCollectionByName")) meta := &MetaTable{ dbName2Meta: map[string]*model.Database{ - util.DefaultDBName: model.NewDefaultDatabase(), + util.DefaultDBName: model.NewDefaultDatabase(nil), }, names: newNameDb(), aliases: newNameDb(), @@ -639,7 +639,7 @@ func TestMetaTable_GetCollectionByName(t *testing.T) { ).Return(&model.Collection{State: pb.CollectionState_CollectionDropped}, nil) meta := &MetaTable{ dbName2Meta: map[string]*model.Database{ - util.DefaultDBName: model.NewDefaultDatabase(), + util.DefaultDBName: model.NewDefaultDatabase(nil), }, names: newNameDb(), aliases: newNameDb(), @@ -669,7 +669,7 @@ func TestMetaTable_GetCollectionByName(t *testing.T) { meta := &MetaTable{ dbName2Meta: map[string]*model.Database{ - util.DefaultDBName: model.NewDefaultDatabase(), + util.DefaultDBName: model.NewDefaultDatabase(nil), }, names: newNameDb(), aliases: newNameDb(), @@ -1206,7 +1206,7 @@ func TestMetaTable_reload(t *testing.T) { catalog.On("ListDatabases", mock.Anything, mock.Anything, - ).Return([]*model.Database{model.NewDefaultDatabase()}, nil) + ).Return([]*model.Database{model.NewDefaultDatabase(nil)}, nil) catalog.On("ListCollections", mock.Anything, mock.Anything, @@ -1508,7 +1508,7 @@ func TestMetaTable_RenameCollection(t *testing.T) { t.Run("new collection name already exist-1", func(t *testing.T) { meta := &MetaTable{ dbName2Meta: map[string]*model.Database{ - util.DefaultDBName: model.NewDefaultDatabase(), + util.DefaultDBName: model.NewDefaultDatabase(nil), }, names: newNameDb(), aliases: newNameDb(), @@ -1535,7 +1535,7 @@ func TestMetaTable_RenameCollection(t *testing.T) { ).Return(nil, errors.New("error mock GetCollectionByID")) meta := &MetaTable{ dbName2Meta: map[string]*model.Database{ - util.DefaultDBName: model.NewDefaultDatabase(), + util.DefaultDBName: model.NewDefaultDatabase(nil), }, catalog: catalog, names: newNameDb(), @@ -1565,7 +1565,7 @@ func TestMetaTable_RenameCollection(t *testing.T) { meta := &MetaTable{ dbName2Meta: map[string]*model.Database{ - util.DefaultDBName: model.NewDefaultDatabase(), + util.DefaultDBName: model.NewDefaultDatabase(nil), }, catalog: catalog, names: newNameDb(), @@ -1592,7 +1592,7 @@ func TestMetaTable_RenameCollection(t *testing.T) { ).Return(nil, merr.WrapErrCollectionNotFound("error")) meta := &MetaTable{ dbName2Meta: map[string]*model.Database{ - util.DefaultDBName: model.NewDefaultDatabase(), + util.DefaultDBName: model.NewDefaultDatabase(nil), "db1": model.NewDatabase(2, "db1", pb.DatabaseState_DatabaseCreated, nil), }, catalog: catalog, @@ -1622,7 +1622,7 @@ func TestMetaTable_RenameCollection(t *testing.T) { ).Return(nil, merr.WrapErrCollectionNotFound("error")) meta := &MetaTable{ dbName2Meta: map[string]*model.Database{ - util.DefaultDBName: model.NewDefaultDatabase(), + util.DefaultDBName: model.NewDefaultDatabase(nil), "db1": model.NewDatabase(2, "db1", pb.DatabaseState_DatabaseCreated, nil), }, catalog: catalog, @@ -1659,7 +1659,7 @@ func TestMetaTable_RenameCollection(t *testing.T) { ).Return(nil, merr.WrapErrCollectionNotFound("error")) meta := &MetaTable{ dbName2Meta: map[string]*model.Database{ - util.DefaultDBName: model.NewDefaultDatabase(), + util.DefaultDBName: model.NewDefaultDatabase(nil), }, catalog: catalog, names: newNameDb(), @@ -1925,7 +1925,7 @@ func TestMetaTable_EmtpyDatabaseName(t *testing.T) { mt := &MetaTable{ names: newNameDb(), dbName2Meta: map[string]*model.Database{ - util.DefaultDBName: model.NewDefaultDatabase(), + util.DefaultDBName: model.NewDefaultDatabase(nil), "db2": model.NewDatabase(2, "db2", pb.DatabaseState_DatabaseCreated, nil), }, collID2Meta: map[typeutil.UniqueID]*model.Collection{ diff --git a/pkg/util/funcutil/func.go b/pkg/util/funcutil/func.go index 82f79560cb525..5ceeb8c036a9a 100644 --- a/pkg/util/funcutil/func.go +++ b/pkg/util/funcutil/func.go @@ -168,6 +168,15 @@ func GetVecFieldIDs(schema *schemapb.CollectionSchema) []int64 { return vecFieldIDs } +func String2KeyValuePair(v string) ([]*commonpb.KeyValuePair, error) { + m := make(map[string]string) + err := json.Unmarshal([]byte(v), &m) + if err != nil { + return nil, err + } + return Map2KeyValuePair(m), nil +} + func Map2KeyValuePair(datas map[string]string) []*commonpb.KeyValuePair { results := make([]*commonpb.KeyValuePair, len(datas)) offset := 0 diff --git a/pkg/util/funcutil/func_test.go b/pkg/util/funcutil/func_test.go index ed9c2555f54a3..0201e5b01ea27 100644 --- a/pkg/util/funcutil/func_test.go +++ b/pkg/util/funcutil/func_test.go @@ -838,3 +838,24 @@ func (s *NumRowsWithSchemaSuite) TestErrorCases() { func TestNumRowsWithSchema(t *testing.T) { suite.Run(t, new(NumRowsWithSchemaSuite)) } + +func TestString2KeyValuePair(t *testing.T) { + t.Run("normal", func(t *testing.T) { + kvs, err := String2KeyValuePair("{\"key\": \"value\"}") + assert.NoError(t, err) + assert.Len(t, kvs, 1) + assert.Equal(t, "key", kvs[0].Key) + assert.Equal(t, "value", kvs[0].Value) + }) + + t.Run("err", func(t *testing.T) { + _, err := String2KeyValuePair("{aa}") + assert.Error(t, err) + }) + + t.Run("empty", func(t *testing.T) { + kvs, err := String2KeyValuePair("{}") + assert.NoError(t, err) + assert.Len(t, kvs, 0) + }) +} diff --git a/pkg/util/paramtable/component_param.go b/pkg/util/paramtable/component_param.go index 0cfb2367a9605..c1355a2fcb0bc 100644 --- a/pkg/util/paramtable/component_param.go +++ b/pkg/util/paramtable/component_param.go @@ -1134,6 +1134,7 @@ type rootCoordConfig struct { MaxGeneralCapacity ParamItem `refreshable:"true"` GracefulStopTimeout ParamItem `refreshable:"true"` UseLockScheduler ParamItem `refreshable:"true"` + DefaultDBProperties ParamItem `refreshable:"false"` } func (p *rootCoordConfig) init(base *BaseTable) { @@ -1217,6 +1218,15 @@ Segments with smaller size than this parameter will not be indexed, and will be Export: false, } p.UseLockScheduler.Init(base.mgr) + + p.DefaultDBProperties = ParamItem{ + Key: "rootCoord.defaultDBProperties", + Version: "2.4.16", + DefaultValue: "{}", + Doc: "default db properties, should be a json string", + Export: false, + } + p.DefaultDBProperties.Init(base.mgr) } // ///////////////////////////////////////////////////////////////////////////// diff --git a/pkg/util/paramtable/component_param_test.go b/pkg/util/paramtable/component_param_test.go index 50376288e1f18..2f5eeeeeda97a 100644 --- a/pkg/util/paramtable/component_param_test.go +++ b/pkg/util/paramtable/component_param_test.go @@ -149,6 +149,10 @@ func TestComponentParam(t *testing.T) { params.Save("rootCoord.gracefulStopTimeout", "100") assert.Equal(t, 100*time.Second, Params.GracefulStopTimeout.GetAsDuration(time.Second)) + assert.Equal(t, "{}", Params.DefaultDBProperties.GetValue()) + params.Save("rootCoord.defaultDBProperties", "{\"key\":\"value\"}") + assert.Equal(t, "{\"key\":\"value\"}", Params.DefaultDBProperties.GetValue()) + SetCreateTime(time.Now()) SetUpdateTime(time.Now()) })