80 lines
1.5 KiB
Go
80 lines
1.5 KiB
Go
package dao
|
|
|
|
import (
|
|
"fmt"
|
|
"github.com/furryboard/spider-scheduler/pkg/conf"
|
|
"github.com/furryboard/spider-scheduler/pkg/dao/model"
|
|
"github.com/furryboard/spider-scheduler/pkg/log"
|
|
"github.com/sirupsen/logrus"
|
|
"gorm.io/driver/postgres"
|
|
"gorm.io/gorm"
|
|
"gorm.io/gorm/logger"
|
|
"sync"
|
|
"time"
|
|
)
|
|
|
|
var (
|
|
once sync.Once
|
|
db *gorm.DB
|
|
)
|
|
|
|
func DB() *gorm.DB {
|
|
once.Do(initDB)
|
|
return db
|
|
|
|
}
|
|
|
|
type LogrusWriter struct {
|
|
log *logrus.Logger
|
|
}
|
|
|
|
func (w *LogrusWriter) Printf(format string, v ...interface{}) {
|
|
logStr := fmt.Sprintf(format, v...)
|
|
w.log.WithField("method", "GORM").Warn(logStr)
|
|
}
|
|
|
|
func NewLogger() *LogrusWriter {
|
|
return &LogrusWriter{log: log.Logger()}
|
|
}
|
|
|
|
func initDB() {
|
|
var err error
|
|
dsn := fmt.Sprintf(
|
|
"host=%s user=%s password=%s dbname=%s port=%d sslmode=disable TimeZone=Asia/Shanghai",
|
|
conf.Conf.Database.Host,
|
|
conf.Conf.Database.User,
|
|
conf.Conf.Database.Password,
|
|
conf.Conf.Database.DB,
|
|
conf.Conf.Database.Port,
|
|
)
|
|
|
|
// >= 1s SQL慢查询
|
|
slowLogger := logger.New(NewLogger(), logger.Config{
|
|
SlowThreshold: time.Second * 1,
|
|
LogLevel: logger.Warn,
|
|
})
|
|
db, err = gorm.Open(postgres.Open(dsn), &gorm.Config{
|
|
Logger: slowLogger,
|
|
})
|
|
if err != nil {
|
|
log.Logger().Fatal("创建数据库连接失败：", err)
|
|
}
|
|
|
|
autoMigrate()
|
|
}
|
|
|
|
func autoMigrate() {
|
|
var err error
|
|
err = db.AutoMigrate(&model.Changelog{})
|
|
err = db.AutoMigrate(&model.Fans{})
|
|
err = db.AutoMigrate(&model.Metadata{})
|
|
err = db.AutoMigrate(&model.LiveRoom{})
|
|
if err != nil {
|
|
return
|
|
}
|
|
|
|
if err != nil {
|
|
log.Logger().Fatal("关联数据表失败：", err)
|
|
}
|
|
}
|