package orm import ( "database/sql" "fmt" "os" "reflect" "sync" "time" ) 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 MaxIdleConns int MaxOpenConns int DB *sql.DB DbBaser dbBaser TZ *time.Location Engine string } // Setting the database connect params. Use the database driver self dataSource args. func RegisterDataBase(aliasName, driverName, dataSource string, params ...int) { al := new(alias) al.Name = aliasName al.DriverName = driverName al.DataSource = dataSource 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(aliasName, al) == false { err = fmt.Errorf("db name `%s` already registered, cannot reuse", aliasName) goto end } al.DB, err = sql.Open(driverName, dataSource) if err != nil { err = fmt.Errorf("register db `%s`, %s", aliasName, err.Error()) goto end } // 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" { tz = "" row = al.DB.QueryRow("SELECT @@system_time_zone") row.Scan(&tz) t, err := time.Parse("MST", tz) if err == nil { al.TZ = t.Location() } } else { t, err := time.Parse("-07:00", tz) if err == nil { al.TZ = t.Location() } } // get default engine from current database row = al.DB.QueryRow("SELECT ENGINE, TRANSACTIONS FROM information_schema.engines WHERE SUPPORT = 'DEFAULT'") var engine string var tx bool row.Scan(&engine, &tx) if engine != "" { al.Engine = engine } else { engine = "INNODB" } 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 } } for i, v := range params { switch i { case 0: SetMaxIdleConns(al.Name, v) case 1: SetMaxOpenConns(al.Name, v) } } err = al.DB.Ping() if err != nil { err = fmt.Errorf("register db `%s`, %s", aliasName, err.Error()) goto end } end: if err != nil { fmt.Println(err.Error()) os.Exit(2) } } // Register a database driver use specify driver name, this can be definition the driver is which database type. 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) } } } // Change the database default used timezone func SetDataBaseTZ(aliasName string, tz *time.Location) { if al, ok := dataBaseCache.get(aliasName); ok { al.TZ = tz } else { fmt.Sprintf("DataBase name `%s` not registered\n", aliasName) os.Exit(2) } } // 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.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 // for tip go 1.2 if fun := reflect.ValueOf(al.DB).MethodByName("SetMaxOpenConns"); fun.IsValid() { fun.Call([]reflect.Value{reflect.ValueOf(maxOpenConns)}) } }