diff --git a/pkg/client/orm/do_nothing_orm.go b/pkg/client/orm/do_nothing_orm.go index e27e7f3a..07c7fd74 100644 --- a/pkg/client/orm/do_nothing_orm.go +++ b/pkg/client/orm/do_nothing_orm.go @@ -30,6 +30,18 @@ var _ Ormer = new(DoNothingOrm) type DoNothingOrm struct { } +func (d *DoNothingOrm) RegisterModels(models ...interface{}) (err error) { + return nil +} + +func (d *DoNothingOrm) RegisterModelsWithPrefix(prefix string, models ...interface{}) (err error) { + return nil +} + +func (d *DoNothingOrm) RegisterModelsWithSuffix(suffix string, models ...interface{}) (err error) { + return nil +} + func (d *DoNothingOrm) Read(md interface{}, cols ...string) error { return nil } diff --git a/pkg/client/orm/filter_orm_decorator.go b/pkg/client/orm/filter_orm_decorator.go index d0c5c537..095c8485 100644 --- a/pkg/client/orm/filter_orm_decorator.go +++ b/pkg/client/orm/filter_orm_decorator.go @@ -32,6 +32,7 @@ var _ TxOrmer = new(filterOrmDecorator) type filterOrmDecorator struct { ormer + modelRegister TxBeginner TxCommitter @@ -42,6 +43,18 @@ type filterOrmDecorator struct { txName string } +func (f *filterOrmDecorator) RegisterModels(models ...interface{}) (err error) { + return f.modelRegister.RegisterModels(models...) +} + +func (f *filterOrmDecorator) RegisterModelsWithPrefix(prefix string, models ...interface{}) (err error) { + return f.modelRegister.RegisterModelsWithPrefix(prefix, models...) +} + +func (f *filterOrmDecorator) RegisterModelsWithSuffix(suffix string, models ...interface{}) (err error) { + return f.modelRegister.RegisterModelsWithSuffix(suffix, models...) +} + func NewFilterOrmDecorator(delegate Ormer, filterChains ...FilterChain) Ormer { res := &filterOrmDecorator{ ormer: delegate, diff --git a/pkg/client/orm/models.go b/pkg/client/orm/models.go index a7de10f7..97faa00a 100644 --- a/pkg/client/orm/models.go +++ b/pkg/client/orm/models.go @@ -39,6 +39,15 @@ var ( } ) +type modelRegister interface { + //RegisterModels register models without prefix or suffix + RegisterModels(models ...interface{}) (err error) + //RegisterModelsWithPrefix register models with prefix + RegisterModelsWithPrefix(prefix string, models ...interface{}) (err error) + //RegisterModelsWithSuffix register models with suffix + RegisterModelsWithSuffix(suffix string, models ...interface{}) (err error) +} + // model info collection type _modelCache struct { sync.RWMutex // only used outsite for bootStrap @@ -48,6 +57,20 @@ type _modelCache struct { done bool } +var _ modelRegister = new(_modelCache) + +func (mc *_modelCache) RegisterModels(models ...interface{}) (err error) { + return mc.register(``, true, models...) +} + +func (mc *_modelCache) RegisterModelsWithPrefix(prefix string, models ...interface{}) (err error) { + return mc.register(prefix, true, models...) +} + +func (mc *_modelCache) RegisterModelsWithSuffix(suffix string, models ...interface{}) (err error) { + return mc.register(suffix, false, models...) +} + // get all model info func (mc *_modelCache) all() map[string]*modelInfo { m := make(map[string]*modelInfo, len(mc.cache)) diff --git a/pkg/client/orm/models_boot.go b/pkg/client/orm/models_boot.go index 407cf536..9a0ce893 100644 --- a/pkg/client/orm/models_boot.go +++ b/pkg/client/orm/models_boot.go @@ -14,15 +14,8 @@ package orm -import ( - "fmt" -) - // RegisterModel register models func RegisterModel(models ...interface{}) { - if modelCache.done { - panic(fmt.Errorf("RegisterModel must be run before BootStrap")) - } RegisterModelWithPrefix("", models...) } diff --git a/pkg/client/orm/orm.go b/pkg/client/orm/orm.go index 634b1892..d82f7e05 100644 --- a/pkg/client/orm/orm.go +++ b/pkg/client/orm/orm.go @@ -498,10 +498,23 @@ func (o *ormBase) DBStats() *sql.DBStats { type orm struct { ormBase + modelRegister } var _ Ormer = new(orm) +func (o *orm) RegisterModels(models ...interface{}) (err error) { + return o.modelRegister.RegisterModels(models) +} + +func (o *orm) RegisterModelsWithPrefix(prefix string, models ...interface{}) (err error) { + return o.modelRegister.RegisterModelsWithPrefix(prefix, models...) +} + +func (o *orm) RegisterModelsWithSuffix(suffix string, models ...interface{}) (err error) { + return o.modelRegister.RegisterModelsWithSuffix(suffix, models...) +} + func (o *orm) Begin() (TxOrmer, error) { return o.BeginWithCtx(context.Background()) } diff --git a/pkg/client/orm/types.go b/pkg/client/orm/types.go index eb34e759..584f0f8a 100644 --- a/pkg/client/orm/types.go +++ b/pkg/client/orm/types.go @@ -214,6 +214,7 @@ type ormer interface { type Ormer interface { ormer + modelRegister TxBeginner }