Refactor RegisterDatabase

pkg/orm/constant.go

@@ -0,0 +1,21 @@
+// Copyright 2020 beego-dev
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package orm
+
+const (
+	MaxIdleConnsKey    = "MaxIdleConns"
+	MaxOpenConnsKey    = "MaxOpenConns"
+	ConnMaxLifetimeKey = "ConnMaxLifetime"
+)

pkg/orm/db_alias.go

@@ -18,10 +18,12 @@ import (
 	"context"
 	"database/sql"
 	"fmt"
-	lru "github.com/hashicorp/golang-lru"
-	"reflect"
 	"sync"
 	"time"
+
+	lru "github.com/hashicorp/golang-lru"
+
+	"github.com/astaxie/beego/pkg/common"
 )
 
 // DriverType database driver constant int.
@@ -63,7 +65,7 @@ var (
 		"tidb":     DRTiDB,
 		"oracle":   DROracle,
 		"oci8":     DROracle, // github.com/mattn/go-oci8
-		"ora":      DROracle, //https://github.com/rana/ora
+		"ora":      DROracle, // https://github.com/rana/ora
 	}
 	dbBasers = map[DriverType]dbBaser{
 		DRMySQL:    newdbBaseMysql(),
@@ -122,7 +124,7 @@ func (d *DB) BeginTx(ctx context.Context, opts *sql.TxOptions) (*sql.Tx, error)
 	return d.DB.BeginTx(ctx, opts)
 }
 
-//su must call release to release *sql.Stmt after using
+// su must call release to release *sql.Stmt after using
 func (d *DB) getStmtDecorator(query string) (*stmtDecorator, error) {
 	d.RLock()
 	c, ok := d.stmtDecorators.Get(query)
@@ -274,16 +276,17 @@ func (t *TxDB) QueryRowContext(ctx context.Context, query string, args ...interf
 }
 
 type alias struct {
-	Name         string
-	Driver       DriverType
-	DriverName   string
-	DataSource   string
-	MaxIdleConns int
-	MaxOpenConns int
-	DB           *DB
-	DbBaser      dbBaser
-	TZ           *time.Location
-	Engine       string
+	Name            string
+	Driver          DriverType
+	DriverName      string
+	DataSource      string
+	MaxIdleConns    int
+	MaxOpenConns    int
+	ConnMaxLifetime time.Duration
+	DB              *DB
+	DbBaser         dbBaser
+	TZ              *time.Location
+	Engine          string
 }
 
 func detectTZ(al *alias) {
@@ -378,13 +381,15 @@ func AddAliasWthDB(aliasName, driverName string, db *sql.DB) error {
 }
 
 // RegisterDataBase Setting the database connect params. Use the database driver self dataSource args.
-func RegisterDataBase(aliasName, driverName, dataSource string, params ...int) error {
+func RegisterDataBase(aliasName, driverName, dataSource string, params ...common.KV) error {
 	var (
 		err error
 		db  *sql.DB
 		al  *alias
 	)
 
+	kvs := common.NewKVs(params...)
+
 	db, err = sql.Open(driverName, dataSource)
 	if err != nil {
 		err = fmt.Errorf("register db `%s`, %s", aliasName, err.Error())
@@ -400,14 +405,13 @@ func RegisterDataBase(aliasName, driverName, dataSource string, params ...int) e
 
 	detectTZ(al)
 
-	for i, v := range params {
-		switch i {
-		case 0:
-			SetMaxIdleConns(al.Name, v)
-		case 1:
-			SetMaxOpenConns(al.Name, v)
-		}
-	}
+	kvs.IfContains(MaxIdleConnsKey, func(value interface{}) {
+		SetMaxIdleConns(al.Name, value.(int))
+	}).IfContains(MaxOpenConnsKey, func(value interface{}) {
+		SetMaxOpenConns(al.Name, value.(int))
+	}).IfContains(ConnMaxLifetimeKey, func(value interface{}) {
+		SetConnMaxLifetime(al.Name, value.(time.Duration))
+	})
 
 end:
 	if err != nil {
@@ -454,10 +458,12 @@ func SetMaxOpenConns(aliasName string, maxOpenConns int) {
 	al := getDbAlias(aliasName)
 	al.MaxOpenConns = maxOpenConns
 	al.DB.DB.SetMaxOpenConns(maxOpenConns)
-	// for tip go 1.2
-	if fun := reflect.ValueOf(al.DB).MethodByName("SetMaxOpenConns"); fun.IsValid() {
-		fun.Call([]reflect.Value{reflect.ValueOf(maxOpenConns)})
-	}
+}
+
+func SetConnMaxLifetime(aliasName string, lifeTime time.Duration) {
+	al := getDbAlias(aliasName)
+	al.ConnMaxLifetime = lifeTime
+	al.DB.DB.SetConnMaxLifetime(lifeTime)
 }
 
 // GetDB Get *sql.DB from registered database by db alias name.
@@ -477,7 +483,7 @@ func GetDB(aliasNames ...string) (*sql.DB, error) {
 }
 
 type stmtDecorator struct {
-	wg sync.WaitGroup
+	wg   sync.WaitGroup
 	stmt *sql.Stmt
 }
 
@@ -497,7 +503,7 @@ func (s *stmtDecorator) release() {
 	s.wg.Done()
 }
 
-//garbage recycle for stmt
+// garbage recycle for stmt
 func (s *stmtDecorator) destroy() {
 	go func() {
 		s.wg.Wait()

pkg/orm/db_alias_test.go

@@ -0,0 +1,44 @@
+// Copyright 2020 beego-dev
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package orm
+
+import (
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+
+	"github.com/astaxie/beego/pkg/common"
+)
+
+func TestRegisterDataBase(t *testing.T) {
+	err := RegisterDataBase("test-params", DBARGS.Driver, DBARGS.Source, common.KV{
+		Key:   MaxIdleConnsKey,
+		Value: 20,
+	}, common.KV{
+		Key:   MaxOpenConnsKey,
+		Value: 300,
+	}, common.KV{
+		Key:   ConnMaxLifetimeKey,
+		Value: time.Minute,
+	})
+	assert.Nil(t, err)
+
+	al := getDbAlias("test-params")
+	assert.NotNil(t, al)
+	assert.Equal(t, al.MaxIdleConns, 20)
+	assert.Equal(t, al.MaxOpenConns, 300)
+	assert.Equal(t, al.ConnMaxLifetime, time.Minute)
+}

pkg/orm/models_test.go

@@ -27,6 +27,8 @@ import (
 	_ "github.com/mattn/go-sqlite3"
 	// As tidb can't use go get, so disable the tidb testing now
 	// _ "github.com/pingcap/tidb"
+
+	"github.com/astaxie/beego/pkg/common"
 )
 
 // A slice string field.
@@ -487,7 +489,10 @@ func init() {
 		os.Exit(2)
 	}
 
-	err := RegisterDataBase("default", DBARGS.Driver, DBARGS.Source, 20)
+	err := RegisterDataBase("default", DBARGS.Driver, DBARGS.Source, common.KV{
+		Key:MaxIdleConnsKey,
+		Value:20,
+	})
 
 	if err != nil{
 		panic(fmt.Sprintf("can not register database: %v", err))