From 7574b91760309df810c6b506ff1d9fda877736d4 Mon Sep 17 00:00:00 2001
From: Anker Jam <jianzhiyao020@qq.com>
Date: Wed, 2 Sep 2020 00:26:25 +0800
Subject: [PATCH] add type modelRegister interface into Ormer

---
 pkg/client/orm/do_nothing_orm.go       | 12 ++++++++++++
 pkg/client/orm/filter_orm_decorator.go | 13 +++++++++++++
 pkg/client/orm/models.go               | 23 +++++++++++++++++++++++
 pkg/client/orm/models_boot.go          |  7 -------
 pkg/client/orm/orm.go                  | 13 +++++++++++++
 pkg/client/orm/types.go                |  1 +
 6 files changed, 62 insertions(+), 7 deletions(-)

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
 }