147 lines
3.8 KiB
Go
147 lines
3.8 KiB
Go
package common
|
|
|
|
import (
|
|
"database/sql"
|
|
"fmt"
|
|
"regexp"
|
|
"strings"
|
|
"time"
|
|
|
|
"github.com/sirupsen/logrus"
|
|
"gorm.io/driver/mysql"
|
|
"gorm.io/gorm"
|
|
)
|
|
|
|
type M = map[string]interface{}
|
|
type MS = map[string]string
|
|
|
|
type ModelBase struct {
|
|
ID uint
|
|
CreateTime *time.Time `gorm:"autoCreateTime"`
|
|
UpdateTime *time.Time `gorm:"autoUpdateTime"`
|
|
}
|
|
|
|
var dbs = map[string]*DB{}
|
|
|
|
type DB struct {
|
|
*gorm.DB
|
|
}
|
|
|
|
func (m *DB) Must() *DB {
|
|
db := m.InstanceSet("ivy.must", true)
|
|
return &DB{DB: db}
|
|
}
|
|
|
|
func (m *DB) NoMust() *DB {
|
|
db := m.InstanceSet("ivy.must", false)
|
|
return &DB{DB: db}
|
|
}
|
|
|
|
func (m *DB) ToSqlDB() *sql.DB {
|
|
d, err := m.DB.DB()
|
|
E2P(err)
|
|
return d
|
|
}
|
|
|
|
type tracePlugin struct{}
|
|
|
|
func (op *tracePlugin) Name() string {
|
|
return "tracePlugin"
|
|
}
|
|
|
|
func (op *tracePlugin) Initialize(db *gorm.DB) (err error) {
|
|
before := func(db *gorm.DB) {
|
|
db.InstanceSet("ivy.startTime", time.Now())
|
|
}
|
|
|
|
after := func(db *gorm.DB) {
|
|
_ts, _ := db.InstanceGet("ivy.startTime")
|
|
sql := db.Dialector.Explain(db.Statement.SQL.String(), db.Statement.Vars...)
|
|
logrus.Printf("used: %d ms affected: %d sql is: %s", time.Since(_ts.(time.Time)).Milliseconds(), db.RowsAffected, sql)
|
|
if v, ok := db.InstanceGet("ivy.must"); ok && v.(bool) {
|
|
if db.Error != nil && db.Error != gorm.ErrRecordNotFound {
|
|
panic(db.Error)
|
|
}
|
|
}
|
|
}
|
|
|
|
beforeName := "cb_before"
|
|
afterName := "cb_after"
|
|
|
|
logrus.Printf("installing db plugin: %s", op.Name())
|
|
// 开始前
|
|
_ = db.Callback().Create().Before("gorm:before_create").Register(beforeName, before)
|
|
_ = db.Callback().Query().Before("gorm:query").Register(beforeName, before)
|
|
_ = db.Callback().Delete().Before("gorm:before_delete").Register(beforeName, before)
|
|
_ = db.Callback().Update().Before("gorm:setup_reflect_value").Register(beforeName, before)
|
|
_ = db.Callback().Row().Before("gorm:row").Register(beforeName, before)
|
|
_ = db.Callback().Raw().Before("gorm:raw").Register(beforeName, before)
|
|
|
|
// 结束后
|
|
_ = db.Callback().Create().After("gorm:after_create").Register(afterName, after)
|
|
_ = db.Callback().Query().After("gorm:after_query").Register(afterName, after)
|
|
_ = db.Callback().Delete().After("gorm:after_delete").Register(afterName, after)
|
|
_ = db.Callback().Update().After("gorm:after_update").Register(afterName, after)
|
|
_ = db.Callback().Row().After("gorm:row").Register(afterName, after)
|
|
_ = db.Callback().Raw().After("gorm:raw").Register(afterName, after)
|
|
return
|
|
}
|
|
|
|
func GetDsn(conf map[string]string) string {
|
|
if IsDocker() {
|
|
conf["host"] = strings.Replace(conf["host"], "localhost", "host.docker.internal", 1)
|
|
}
|
|
return fmt.Sprintf("%s:%s@tcp(%s:%s)/%s?charset=utf8mb4&parseTime=true&loc=Local", conf["user"], conf["password"], conf["host"], conf["port"], conf["database"])
|
|
}
|
|
|
|
func ReplaceDsnPassword(dsn string) string {
|
|
reg := regexp.MustCompile(`:.*@`)
|
|
return reg.ReplaceAllString(dsn, ":****@")
|
|
}
|
|
|
|
func DbGet(conf map[string]string) *DB {
|
|
dsn := GetDsn(conf)
|
|
if dbs[dsn] == nil {
|
|
logrus.Printf("connecting %s", ReplaceDsnPassword(dsn))
|
|
db1, err := gorm.Open(mysql.Open(dsn), &gorm.Config{
|
|
SkipDefaultTransaction: true,
|
|
})
|
|
E2P(err)
|
|
db1.Use(&tracePlugin{})
|
|
dbs[dsn] = &DB{DB: db1}
|
|
}
|
|
return dbs[dsn]
|
|
}
|
|
|
|
func SqlDB2DB(sdb *sql.DB) *DB {
|
|
db, err := gorm.Open(mysql.New(mysql.Config{
|
|
Conn: sdb,
|
|
}), &gorm.Config{})
|
|
E2P(err)
|
|
db.Use(&tracePlugin{})
|
|
return &DB{DB: db}
|
|
}
|
|
|
|
type MyConn struct {
|
|
Conn *sql.DB
|
|
Dsn string
|
|
}
|
|
|
|
func (conn *MyConn) Close() {
|
|
logrus.Printf("closing alone mysql: %s", ReplaceDsnPassword(conn.Dsn))
|
|
conn.Conn.Close()
|
|
}
|
|
|
|
func DbAlone(conf map[string]string) (*DB, *MyConn) {
|
|
dsn := GetDsn(conf)
|
|
logrus.Printf("opening alone mysql: %s", ReplaceDsnPassword(dsn))
|
|
mdb, err := sql.Open("mysql", dsn)
|
|
E2P(err)
|
|
gormDB, err := gorm.Open(mysql.New(mysql.Config{
|
|
Conn: mdb,
|
|
}), &gorm.Config{})
|
|
E2P(err)
|
|
gormDB.Use(&tracePlugin{})
|
|
return &DB{DB: gormDB}, &MyConn{Conn: mdb, Dsn: dsn}
|
|
}
|