Beego/orm/db_alias.go

package orm

import (
	"database/sql"
	"fmt"
	"os"
	"sync"
	"time"
)

const defaultMaxIdle = 30

type DriverType int

const (
	_ DriverType = iota
	DR_MySQL
	DR_Sqlite
	DR_Oracle
	DR_Postgres
)

type driver string

func (d driver) Type() DriverType {
	a, _ := dataBaseCache.get(string(d))
	return a.Driver
}

func (d driver) Name() string {
	return string(d)
}

var _ Driver = new(driver)

var (
	dataBaseCache = &_dbCache{cache: make(map[string]*alias)}
	drivers       = map[string]DriverType{
		"mysql":    DR_MySQL,
		"postgres": DR_Postgres,
		"sqlite3":  DR_Sqlite,
	}
	dbBasers = map[DriverType]dbBaser{
		DR_MySQL:    newdbBaseMysql(),
		DR_Sqlite:   newdbBaseSqlite(),
		DR_Oracle:   newdbBaseMysql(),
		DR_Postgres: newdbBasePostgres(),
	}
)

type _dbCache struct {
	mux   sync.RWMutex
	cache map[string]*alias
}

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
}

func (ac *_dbCache) get(name string) (al *alias, ok bool) {
	ac.mux.RLock()
	defer ac.mux.RUnlock()
	al, ok = ac.cache[name]
	return
}

func (ac *_dbCache) getDefault() (al *alias) {
	al, _ = ac.get("default")
	return
}

type alias struct {
	Name       string
	Driver     DriverType
	DriverName string
	DataSource string
	MaxIdle    int
	DB         *sql.DB
	DbBaser    dbBaser
	TZ         *time.Location
}

func RegisterDataBase(name, driverName, dataSource string, maxIdle int) {
	if maxIdle <= 0 {
		maxIdle = defaultMaxIdle
	}

	al := new(alias)
	al.Name = name
	al.DriverName = driverName
	al.DataSource = dataSource
	al.MaxIdle = maxIdle

	var (
		err error
	)

	if dr, ok := drivers[driverName]; ok {
		al.DbBaser = dbBasers[dr]
		al.Driver = dr
	} else {
		err = fmt.Errorf("driver name `%s` have not registered", driverName)
		goto end
	}

	if dataBaseCache.add(name, al) == false {
		err = fmt.Errorf("db name `%s` already registered, cannot reuse", name)
		goto end
	}

	al.DB, err = sql.Open(driverName, dataSource)
	if err != nil {
		err = fmt.Errorf("register db `%s`, %s", name, err.Error())
		goto end
	}

	al.DB.SetMaxIdleConns(al.MaxIdle)

	// orm timezone system match database
	// default use Local
	al.TZ = time.Local

	switch al.Driver {
	case DR_MySQL:
		row := al.DB.QueryRow("SELECT @@session.time_zone")
		var tz string
		row.Scan(&tz)
		if tz != "SYSTEM" {
			t, err := time.Parse("-07:00", tz)
			if err == nil {
				al.TZ = t.Location()
			}
		}
	case DR_Sqlite:
		al.TZ = time.UTC
	case DR_Postgres:
		row := al.DB.QueryRow("SELECT current_setting('TIMEZONE')")
		var tz string
		row.Scan(&tz)
		loc, err := time.LoadLocation(tz)
		if err == nil {
			al.TZ = loc
		}
	}

	err = al.DB.Ping()
	if err != nil {
		err = fmt.Errorf("register db `%s`, %s", name, err.Error())
		goto end
	}

end:
	if err != nil {
		fmt.Println(err.Error())
		os.Exit(2)
	}
}

func RegisterDriver(driverName string, typ DriverType) {
	if t, ok := drivers[driverName]; ok == false {
		drivers[driverName] = typ
	} else {
		if t != typ {
			fmt.Sprintf("driverName `%s` db driver already registered and is other type\n", driverName)
			os.Exit(2)
		}
	}
}

func SetDataBaseTZ(name string, tz *time.Location) {
	if al, ok := dataBaseCache.get(name); ok {
		al.TZ = tz
	} else {
		fmt.Sprintf("DataBase name `%s` not registered\n", name)
		os.Exit(2)
	}
}