rebuild gorm (#53)

example/gorm/main.go

@@ -5,10 +5,11 @@ import (
 
 	"github.com/sunmi-OS/gocore/gorm"
 	"github.com/sunmi-OS/gocore/viper"
+	"github.com/sunmi-OS/gocore/xlog"
 )
 
 type Machine struct {
-	MId int64  `gorm:"column:mId"`
+	Mid int64  `gorm:"column:mid"`
 	Msn string `gorm:"column:msn"`
 }
 
@@ -20,16 +21,24 @@ func main() {
 	// 指定配置文件所在的目录和文件名称
 	viper.NewConfig("config", "conf")
 
-	gorm.NewDB("dbDefault")
+	gorm.NewDB("a")
+	gorm.NewDB("b")
+	gorm.NewDB("c")
 
-	MC := []Machine{}
+	client := gorm.Gorm()
+	err := client.NewOrUpdateDB("d")
+	if err != nil {
+		xlog.Errorf("NewOrUpdateDB(%s),error:%+v", "d", err)
+	}
 
-	err := gorm.GetORM().Where("msn =  ?", "7102V04115500128").Find(&MC).Error
+	var MC []Machine
+
+	err = client.GetORM().Where("msn =  ?", "7102V04115500128").Find(&MC).Error
 
 	if err != nil {
 		fmt.Println(err.Error())
 	}
-
 	fmt.Println(MC)
 
+	client.Close()
 }

gorm/gorm.go

@@ -12,31 +12,28 @@ import (
 	"github.com/sunmi-OS/gocore/xlog"
 )
 
-var (
-	Gorm        sync.Map
-	defaultName = "dbDefault"
-)
+type Client struct {
+	maps          sync.Map
+	defaultDbName string
+}
 
-var (
-	// ErrRecordNotFound record not found error, happens when haven't find any matched data when looking up with a struct
-	ErrRecordNotFound = gorm.ErrRecordNotFound
-	// ErrInvalidSQL invalid SQL error, happens when you passed invalid SQL
-	ErrInvalidSQL = gorm.ErrInvalidSQL
-	// ErrInvalidTransaction invalid transaction when you are trying to `Commit` or `Rollback`
-	ErrInvalidTransaction = gorm.ErrInvalidTransaction
-	// ErrCantStartTransaction can't start transaction when you are trying to start one with `Begin`
-	ErrCantStartTransaction = gorm.ErrCantStartTransaction
-	// ErrUnaddressable unaddressable value
-	ErrUnaddressable = gorm.ErrUnaddressable
-)
+var _Gorm *Client
+
+func Gorm() *Client {
+	return _Gorm
+}
 
 // 初始化Gorm
 func NewDB(dbname string) {
 	var (
 		orm *gorm.DB
 		err error
 	)
+	if _Gorm == nil {
+		_Gorm = &Client{defaultDbName: defaultName}
+	}
 
+	// openORM
 	err = retry.Retry(func() error {
 		orm, err = openORM(dbname)
 		if err != nil {
@@ -49,16 +46,17 @@ func NewDB(dbname string) {
 		panic(err)
 	}
 
-	Gorm.Store(dbname, orm)
+	// store db client
+	_Gorm.maps.Store(dbname, orm)
 }
 
-// 设置获取db的默认值
-func SetDefaultName(dbname string) {
-	defaultName = dbname
+// SetDefaultName 设置默认DB Name
+func (c *Client) SetDefaultName(dbName string) {
+	c.defaultDbName = dbName
 }
 
-// 初始化Gorm
-func UpdateDB(dbname string) error {
+// NewOrUpdateDB 初始化或更新Gorm
+func (c *Client) NewOrUpdateDB(dbname string) error {
 	var (
 		orm *gorm.DB
 		err error
@@ -78,11 +76,11 @@ func UpdateDB(dbname string) error {
 	}
 
 	// second: load gorm client
-	v, _ := Gorm.Load(dbname)
+	v, _ := c.maps.Load(dbname)
 
 	// third: delete old gorm client and store the new gorm client
-	Gorm.Delete(dbname)
-	Gorm.Store(dbname, orm)
+	c.maps.Delete(dbname)
+	c.maps.Store(dbname, orm)
 
 	// fourth: if old client is not nil, delete and close connection
 	if v != nil {
@@ -91,30 +89,30 @@ func UpdateDB(dbname string) error {
 	return nil
 }
 
-// Deprecated
-// 通过名称获取Gorm实例
-func GetORMByName(dbname string) *gorm.DB {
-	v, ok := Gorm.Load(dbname)
-	if ok {
-		return v.(*gorm.DB)
-	}
-	return nil
-}
-
 // GetORM 获取默认的Gorm实例
 // 目前仅支持 不传 或者仅传一个 dbname
-func GetORM(dbname ...string) *gorm.DB {
-	name := defaultName
+func (c *Client) GetORM(dbname ...string) *gorm.DB {
+	name := c.defaultDbName
 	if len(dbname) == 1 {
 		name = dbname[0]
 	}
-	v, ok := Gorm.Load(name)
+
+	v, ok := c.maps.Load(name)
 	if ok {
 		return v.(*gorm.DB)
 	}
 	return nil
 }
 
+func (c *Client) Close() {
+	c.maps.Range(func(dbName, orm interface{}) bool {
+		xlog.Warnf("close db %s", dbName)
+		c.maps.Delete(dbName)
+		orm.(*gorm.DB).Close()
+		return true
+	})
+}
+
 func openORM(dbname string) (*gorm.DB, error) {
 	//默认配置
 	viper.C.SetDefault(dbname, map[string]interface{}{

gorm/gorm_test.go

@@ -0,0 +1,16 @@
+package gorm
+
+import (
+	"testing"
+)
+
+func TestNewDB(t *testing.T) {
+	//NewDB("a")
+	//NewDB("b")
+	//NewDB("c")
+
+	// only example
+	//Gorm().GetORM().Select("id").Where("id = ?", 1).First(nil)
+
+	//Gorm().Close()
+}

gorm/model.go

@@ -0,0 +1,20 @@
+package gorm
+
+import "github.com/jinzhu/gorm"
+
+const (
+	defaultName = "dbDefault"
+)
+
+var (
+	// ErrRecordNotFound record not found error, happens when haven't find any matched data when looking up with a struct
+	ErrRecordNotFound = gorm.ErrRecordNotFound
+	// ErrInvalidSQL invalid SQL error, happens when you passed invalid SQL
+	ErrInvalidSQL = gorm.ErrInvalidSQL
+	// ErrInvalidTransaction invalid transaction when you are trying to `Commit` or `Rollback`
+	ErrInvalidTransaction = gorm.ErrInvalidTransaction
+	// ErrCantStartTransaction can't start transaction when you are trying to start one with `Begin`
+	ErrCantStartTransaction = gorm.ErrCantStartTransaction
+	// ErrUnaddressable unaddressable value
+	ErrUnaddressable = gorm.ErrUnaddressable
+)

retry/retry.go

@@ -2,16 +2,13 @@ package retry
 
 import (
 	"time"
-
-	"github.com/sunmi-OS/gocore/xlog"
 )
 
 // Retry 重试 func 最大次数，间隔
 func Retry(callback func() error, maxRetries int, interval time.Duration) error {
 	var err error
 	for i := 1; i <= maxRetries; i++ {
 		if err = callback(); err != nil {
-			xlog.Warnf("Retry(%d) error(%+v)", i, err)
 			time.Sleep(interval)
 			continue
 		}