增加了连接缓存,参考beego orm
This commit is contained in:
		
							parent
							
								
									3bd5131132
								
							
						
					
					
						commit
						c2f7bf5323
					
				
							
								
								
									
										286
									
								
								pool.go
									
									
									
									
									
										Normal file
									
								
							
							
						
						
									
										286
									
								
								pool.go
									
									
									
									
									
										Normal file
									
								
							@ -0,0 +1,286 @@
 | 
			
		||||
package gorm
 | 
			
		||||
 | 
			
		||||
import (
 | 
			
		||||
"fmt"
 | 
			
		||||
"time"
 | 
			
		||||
"sync"
 | 
			
		||||
 | 
			
		||||
)
 | 
			
		||||
 | 
			
		||||
 | 
			
		||||
// database driver constant int.
 | 
			
		||||
type DriverType int
 | 
			
		||||
 | 
			
		||||
const (
 | 
			
		||||
_           DriverType = iota // int enum type
 | 
			
		||||
DR_MySQL                      // mysql
 | 
			
		||||
DR_Sqlite                     // sqlite
 | 
			
		||||
DR_Oracle                     // oracle
 | 
			
		||||
DR_Postgres                   // pgsql
 | 
			
		||||
)
 | 
			
		||||
 | 
			
		||||
// database driver string.
 | 
			
		||||
type driverDB string
 | 
			
		||||
 | 
			
		||||
// get type constant int of current driver..
 | 
			
		||||
func (d driverDB) Type() DriverType {
 | 
			
		||||
a, _ := dataBaseCache.get(string(d))
 | 
			
		||||
return a.Driver
 | 
			
		||||
}
 | 
			
		||||
 | 
			
		||||
// get name of current driver
 | 
			
		||||
func (d driverDB) Name() string {
 | 
			
		||||
return string(d)
 | 
			
		||||
}
 | 
			
		||||
 | 
			
		||||
 | 
			
		||||
 | 
			
		||||
var (
 | 
			
		||||
dataBaseCache = &_dbCache{cache: make(map[string]*alias)}
 | 
			
		||||
drivers = map[string]DriverType{
 | 
			
		||||
"mysql":    DR_MySQL,
 | 
			
		||||
"postgres": DR_Postgres,
 | 
			
		||||
"sqlite3":  DR_Sqlite,
 | 
			
		||||
}
 | 
			
		||||
 | 
			
		||||
)
 | 
			
		||||
 | 
			
		||||
// database alias cacher.
 | 
			
		||||
type _dbCache struct {
 | 
			
		||||
mux   sync.RWMutex
 | 
			
		||||
cache map[string]*alias
 | 
			
		||||
}
 | 
			
		||||
 | 
			
		||||
// add database alias with original name.
 | 
			
		||||
func (ac *_dbCache) add(name string, al *alias) (added bool) {
 | 
			
		||||
ac.mux.Lock()
 | 
			
		||||
defer ac.mux.Unlock()
 | 
			
		||||
if _, ok := ac.cache[name]; ok == false {
 | 
			
		||||
ac.cache[name] = al
 | 
			
		||||
added = true
 | 
			
		||||
}
 | 
			
		||||
return
 | 
			
		||||
}
 | 
			
		||||
 | 
			
		||||
// get database alias if cached.
 | 
			
		||||
func (ac *_dbCache) get(name string) (al *alias, ok bool) {
 | 
			
		||||
ac.mux.RLock()
 | 
			
		||||
defer ac.mux.RUnlock()
 | 
			
		||||
al, ok = ac.cache[name]
 | 
			
		||||
return
 | 
			
		||||
}
 | 
			
		||||
 | 
			
		||||
// get default alias.
 | 
			
		||||
func (ac *_dbCache) getDefault() (al *alias) {
 | 
			
		||||
al, _ = ac.get("default")
 | 
			
		||||
return
 | 
			
		||||
}
 | 
			
		||||
 | 
			
		||||
type alias struct {
 | 
			
		||||
Name         string
 | 
			
		||||
Driver       DriverType
 | 
			
		||||
DriverName   string
 | 
			
		||||
DataSource   string
 | 
			
		||||
MaxIdleConns int
 | 
			
		||||
MaxOpenConns int
 | 
			
		||||
DB           *DB
 | 
			
		||||
TZ           *time.Location
 | 
			
		||||
Engine       string
 | 
			
		||||
}
 | 
			
		||||
 | 
			
		||||
 | 
			
		||||
func addAliasWthDB(aliasName, driverName string, db *DB) (*alias, error) {
 | 
			
		||||
al := new(alias)
 | 
			
		||||
al.Name = aliasName
 | 
			
		||||
al.DriverName = driverName
 | 
			
		||||
al.DB = db
 | 
			
		||||
 | 
			
		||||
if dr, ok := drivers[driverName]; ok {
 | 
			
		||||
al.Driver = dr
 | 
			
		||||
} else {
 | 
			
		||||
return nil, fmt.Errorf("driver name `%s` have not registered", driverName)
 | 
			
		||||
}
 | 
			
		||||
 | 
			
		||||
 | 
			
		||||
if dataBaseCache.add(aliasName, al) == false {
 | 
			
		||||
return nil, fmt.Errorf("DataBase alias name `%s` already registered, cannot reuse", aliasName)
 | 
			
		||||
}
 | 
			
		||||
 | 
			
		||||
return al, nil
 | 
			
		||||
}
 | 
			
		||||
 | 
			
		||||
func AddAliasWthDB(aliasName, driverName string, db *DB) error {
 | 
			
		||||
_, err := addAliasWthDB(aliasName, driverName, db)
 | 
			
		||||
return err
 | 
			
		||||
}
 | 
			
		||||
 | 
			
		||||
// Setting the database connect params. Use the database driver self dataSource args.
 | 
			
		||||
func RegisterDataBase(aliasName, driverName, dataSource string, params ...int) error {
 | 
			
		||||
var (
 | 
			
		||||
db *DB
 | 
			
		||||
al  *alias
 | 
			
		||||
)
 | 
			
		||||
 | 
			
		||||
//	db, err = sql.Open(driverName, dataSource)
 | 
			
		||||
_db, err := Open(driverName, dataSource)
 | 
			
		||||
db=&_db
 | 
			
		||||
 | 
			
		||||
 | 
			
		||||
if err != nil {
 | 
			
		||||
err = fmt.Errorf("register db `%s`, %s", aliasName, err.Error())
 | 
			
		||||
goto end
 | 
			
		||||
}
 | 
			
		||||
 | 
			
		||||
al, err = addAliasWthDB(aliasName, driverName, db)
 | 
			
		||||
if err != nil {
 | 
			
		||||
goto end
 | 
			
		||||
}
 | 
			
		||||
 | 
			
		||||
al.DataSource = dataSource
 | 
			
		||||
 | 
			
		||||
 | 
			
		||||
for i, v := range params {
 | 
			
		||||
switch i {
 | 
			
		||||
case 0:
 | 
			
		||||
SetMaxIdleConns(al.Name, v)
 | 
			
		||||
case 1:
 | 
			
		||||
SetMaxOpenConns(al.Name, v)
 | 
			
		||||
}
 | 
			
		||||
}
 | 
			
		||||
 | 
			
		||||
end:
 | 
			
		||||
if err != nil {
 | 
			
		||||
if db != nil {
 | 
			
		||||
db.Close()
 | 
			
		||||
}
 | 
			
		||||
}
 | 
			
		||||
 | 
			
		||||
return err
 | 
			
		||||
}
 | 
			
		||||
 | 
			
		||||
 | 
			
		||||
 | 
			
		||||
// Register a database driver use specify driver name, this can be definition the driver is which database type.
 | 
			
		||||
func RegisterDriver(driverName string, typ DriverType) error {
 | 
			
		||||
if t, ok := drivers[driverName]; ok == false {
 | 
			
		||||
drivers[driverName] = typ
 | 
			
		||||
} else {
 | 
			
		||||
if t != typ {
 | 
			
		||||
return fmt.Errorf("driverName `%s` db driver already registered and is other type\n", driverName)
 | 
			
		||||
}
 | 
			
		||||
}
 | 
			
		||||
return nil
 | 
			
		||||
}
 | 
			
		||||
 | 
			
		||||
 | 
			
		||||
// Change the database default used timezone
 | 
			
		||||
func SetDataBaseTZ(aliasName string, tz *time.Location) error {
 | 
			
		||||
if al, ok := dataBaseCache.get(aliasName); ok {
 | 
			
		||||
al.TZ = tz
 | 
			
		||||
} else {
 | 
			
		||||
return fmt.Errorf("DataBase alias name `%s` not registered\n", aliasName)
 | 
			
		||||
}
 | 
			
		||||
return nil
 | 
			
		||||
}
 | 
			
		||||
func detectTZ(al *alias) {
 | 
			
		||||
// orm timezone system match database
 | 
			
		||||
// default use Local
 | 
			
		||||
al.TZ = time.Local
 | 
			
		||||
 | 
			
		||||
if al.DriverName == "sphinx" {
 | 
			
		||||
return
 | 
			
		||||
}
 | 
			
		||||
 | 
			
		||||
switch al.Driver {
 | 
			
		||||
case DR_MySQL:
 | 
			
		||||
row := al.DB.Raw("SELECT TIMEDIFF(NOW(), UTC_TIMESTAMP)")
 | 
			
		||||
var tz string
 | 
			
		||||
row.Scan(&tz)
 | 
			
		||||
if len(tz) >= 8 {
 | 
			
		||||
if tz[0] != '-' {
 | 
			
		||||
tz = "+" + tz
 | 
			
		||||
}
 | 
			
		||||
t, err := time.Parse("-07:00:00", tz)
 | 
			
		||||
if err == nil {
 | 
			
		||||
al.TZ = t.Location()
 | 
			
		||||
} else {
 | 
			
		||||
fmt.Printf("Detect DB timezone: %s %s\n", tz, err.Error())
 | 
			
		||||
}
 | 
			
		||||
}
 | 
			
		||||
 | 
			
		||||
// get default engine from current database
 | 
			
		||||
row = al.DB.Raw("SELECT ENGINE, TRANSACTIONS FROM information_schema.engines WHERE SUPPORT = 'DEFAULT'")
 | 
			
		||||
 | 
			
		||||
type Egtx struct {
 | 
			
		||||
engine string
 | 
			
		||||
tx bool
 | 
			
		||||
}
 | 
			
		||||
var eg Egtx
 | 
			
		||||
row.Scan(&eg)
 | 
			
		||||
 | 
			
		||||
if eg.engine != "" {
 | 
			
		||||
al.Engine = eg.engine
 | 
			
		||||
} else {
 | 
			
		||||
al.Engine = "INNODB"
 | 
			
		||||
}
 | 
			
		||||
 | 
			
		||||
case DR_Sqlite:
 | 
			
		||||
al.TZ = time.UTC
 | 
			
		||||
 | 
			
		||||
case DR_Postgres:
 | 
			
		||||
row := al.DB.Raw("SELECT current_setting('TIMEZONE')")
 | 
			
		||||
var tz string
 | 
			
		||||
row.Scan(&tz)
 | 
			
		||||
loc, err := time.LoadLocation(tz)
 | 
			
		||||
if err == nil {
 | 
			
		||||
al.TZ = loc
 | 
			
		||||
} else {
 | 
			
		||||
fmt.Printf("Detect DB timezone: %s %s\n", tz, err.Error())
 | 
			
		||||
}
 | 
			
		||||
}
 | 
			
		||||
}
 | 
			
		||||
 | 
			
		||||
 | 
			
		||||
// get table alias.
 | 
			
		||||
func getDbAlias(name string) *alias {
 | 
			
		||||
if al, ok := dataBaseCache.get(name); ok {
 | 
			
		||||
return al
 | 
			
		||||
} else {
 | 
			
		||||
panic(fmt.Errorf("unknown DataBase alias name %s", name))
 | 
			
		||||
}
 | 
			
		||||
}
 | 
			
		||||
 | 
			
		||||
// Change the max idle conns for *sql.DB, use specify database alias name
 | 
			
		||||
func SetMaxIdleConns(aliasName string, maxIdleConns int) {
 | 
			
		||||
al := getDbAlias(aliasName)
 | 
			
		||||
al.MaxIdleConns = maxIdleConns
 | 
			
		||||
al.DB.DB().SetMaxIdleConns(maxIdleConns)
 | 
			
		||||
}
 | 
			
		||||
 | 
			
		||||
// Change the max open conns for *sql.DB, use specify database alias name
 | 
			
		||||
func SetMaxOpenConns(aliasName string, maxOpenConns int) {
 | 
			
		||||
al := getDbAlias(aliasName)
 | 
			
		||||
al.MaxOpenConns = maxOpenConns
 | 
			
		||||
al.DB.DB().SetMaxOpenConns(maxOpenConns)
 | 
			
		||||
 | 
			
		||||
}
 | 
			
		||||
 | 
			
		||||
// Get *sql.DB from registered database by db alias name.
 | 
			
		||||
// Use "default" as alias name if you not set.
 | 
			
		||||
func GetDB(aliasNames ...string) (*DB, error) {
 | 
			
		||||
var name string
 | 
			
		||||
if len(aliasNames) > 0 {
 | 
			
		||||
name = aliasNames[0]
 | 
			
		||||
} else {
 | 
			
		||||
name = "default"
 | 
			
		||||
}
 | 
			
		||||
if al, ok := dataBaseCache.get(name); ok {
 | 
			
		||||
return al.DB, nil
 | 
			
		||||
} else {
 | 
			
		||||
return nil, fmt.Errorf("DataBase of alias name `%s` not found\n", name)
 | 
			
		||||
}
 | 
			
		||||
}
 | 
			
		||||
 | 
			
		||||
 | 
			
		||||
 | 
			
		||||
 | 
			
		||||
		Loading…
	
	
			
			x
			
			
		
	
		Reference in New Issue
	
	Block a user