mc-client-updater-server/pkg/dao/dao.go

package dao

import (
	"fmt"
	"github.com/sirupsen/logrus"
	"gorm.io/driver/postgres"
	"gorm.io/gorm"
	"gorm.io/gorm/logger"
	"mc-client-updater-server/pkg/conf"
	"mc-client-updater-server/pkg/dao/entity"
	"mc-client-updater-server/pkg/log"
	"sync"
	"time"
)

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}
}

var (
	once sync.Once
	db   *gorm.DB
)

func DB() *gorm.DB {
	once.Do(NewDBConn)
	return db
}

func NewDBConn() {
	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)
	}

	migrate()
}

func migrate() {
	var err error
	err = db.AutoMigrate(&entity.User{})
	err = db.AutoMigrate(&entity.Instance{})
	err = db.AutoMigrate(&entity.Update{})
	err = db.AutoMigrate(&entity.Grant{})
	err = db.AutoMigrate(&entity.Token{})
	err = db.AutoMigrate(&entity.Metadata{})

	if err != nil {
		log.Logger.Fatal("关联数据表失败：", err)
	}
}