diff --git a/internal/rootcoord/create_collection_task.go b/internal/rootcoord/create_collection_task.go index 7190efaf21..abe8a7f9ae 100644 --- a/internal/rootcoord/create_collection_task.go +++ b/internal/rootcoord/create_collection_task.go @@ -494,8 +494,8 @@ func (t *createCollectionTask) Prepare(ctx context.Context) error { t.Req.Properties = append(properties, timezoneKV) } - if hookutil.GetEzPropByDBProperties(db.Properties) != nil { - t.Req.Properties = append(t.Req.Properties, hookutil.GetEzPropByDBProperties(db.Properties)) + if ezProps := hookutil.GetEzPropByDBProperties(db.Properties); ezProps != nil { + t.Req.Properties = append(t.Req.Properties, ezProps) } t.header.DbId = db.ID diff --git a/internal/rootcoord/ddl_callbacks_create_database.go b/internal/rootcoord/ddl_callbacks_create_database.go index 2cda570f38..2d777db4e3 100644 --- a/internal/rootcoord/ddl_callbacks_create_database.go +++ b/internal/rootcoord/ddl_callbacks_create_database.go @@ -58,10 +58,16 @@ func (c *Core) broadcastCreateDatabase(ctx context.Context, req *milvuspb.Create if err != nil { return errors.Wrap(err, "failed to tidy database cipher properties") } + tz, exist := funcutil.TryGetAttrByKeyFromRepeatedKV(common.TimezoneKey, properties) if exist && !timestamptz.IsTimezoneValid(tz) { return merr.WrapErrParameterInvalidMsg("unknown or invalid IANA Time Zone ID: %s", tz) } + + if err := hookutil.CreateEZByDBProperties(properties); err != nil { + return errors.Wrap(err, "failed to create ez by db properties") + } + msg := message.NewCreateDatabaseMessageBuilderV2(). WithHeader(&message.CreateDatabaseMessageHeader{ DbName: req.GetDbName(), diff --git a/internal/rootcoord/ddl_callbacks_drop_database.go b/internal/rootcoord/ddl_callbacks_drop_database.go index 5c2b6e79e1..f335ff30af 100644 --- a/internal/rootcoord/ddl_callbacks_drop_database.go +++ b/internal/rootcoord/ddl_callbacks_drop_database.go @@ -25,6 +25,7 @@ import ( "github.com/milvus-io/milvus-proto/go-api/v2/commonpb" "github.com/milvus-io/milvus-proto/go-api/v2/milvuspb" "github.com/milvus-io/milvus/internal/distributed/streaming" + "github.com/milvus-io/milvus/internal/util/hookutil" "github.com/milvus-io/milvus/pkg/v2/streaming/util/message" "github.com/milvus-io/milvus/pkg/v2/streaming/util/message/ce" "github.com/milvus-io/milvus/pkg/v2/util/typeutil" @@ -47,6 +48,11 @@ func (c *Core) broadcastDropDatabase(ctx context.Context, req *milvuspb.DropData return errors.Wrap(err, "failed to get database name") } + // Call back cipher plugin when dropping database succeeded + if err := hookutil.RemoveEZByDBProperties(db.Properties); err != nil { + return errors.Wrap(err, "failed to remove ez by db properties") + } + msg := message.NewDropDatabaseMessageBuilderV2(). WithHeader(&message.DropDatabaseMessageHeader{ DbName: req.GetDbName(), diff --git a/internal/rootcoord/meta_table.go b/internal/rootcoord/meta_table.go index aa47872b86..d10d863aac 100644 --- a/internal/rootcoord/meta_table.go +++ b/internal/rootcoord/meta_table.go @@ -341,6 +341,10 @@ func (mt *MetaTable) createDefaultDb() error { cipherProps := hookutil.GetDBCipherProperties(ezID, defaultRootKey) defaultProperties = append(defaultProperties, cipherProps...) + + if err := hookutil.CreateEZByDBProperties(defaultProperties); err != nil { + return err + } } return mt.createDatabasePrivate(mt.ctx, model.NewDefaultDatabase(defaultProperties), ts) @@ -377,12 +381,7 @@ func (mt *MetaTable) CreateDatabase(ctx context.Context, db *model.Database, ts func (mt *MetaTable) createDatabasePrivate(ctx context.Context, db *model.Database, ts typeutil.Timestamp) error { dbName := db.Name - if err := hookutil.CreateEZByDBProperties(db.Properties); err != nil { - return err - } - if err := mt.catalog.CreateDatabase(ctx, db, ts); err != nil { - hookutil.RemoveEZByDBProperties(db.Properties) // ignore the error since create database failed return err } @@ -444,11 +443,6 @@ func (mt *MetaTable) DropDatabase(ctx context.Context, dbName string, ts typeuti return err } - // Call back cipher plugin when dropping database succeeded - if err := hookutil.RemoveEZByDBProperties(db.Properties); err != nil { - return err - } - mt.names.dropDb(dbName) mt.aliases.dropDb(dbName) delete(mt.dbName2Meta, dbName)