data.go 1.62 KB
Newer Older
彭芳 committed
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17
package data

import (
	"time"

	"gorm.io/driver/mysql"

	"demo/internal/conf"

	"github.com/go-kratos/kratos/v2/log"
	"github.com/google/wire"
	"github.com/sirupsen/logrus"
	"gorm.io/gorm"
	"gorm.io/gorm/logger"
)

// ProviderSet is data providers.
18
var ProviderSet = wire.NewSet(NewData, NewCampaignRepo, NewAdRepo, NewCreativeRepo)
彭芳 committed
19 20 21 22 23 24 25 26 27 28 29

// Data .
type Data struct {
	db *gorm.DB
}

// NewData .
func NewData(c *conf.Data, logger log.Logger) (*Data, func(), error) {
	cleanup := func() {
		log.NewHelper(logger).Info("closing the data resources")
	}
彭芳 committed
30 31 32 33 34 35 36 37
	db, err := newDb(c)
	if err != nil {
		return nil, cleanup, err
	}
	return &Data{db: db}, cleanup, nil
}

func newDb(c *conf.Data) (*gorm.DB, error) {
彭芳 committed
38 39 40 41 42 43 44
	db, err := gorm.Open(mysql.New(mysql.Config{
		DriverName: c.Database.Driver,
		DSN:        c.Database.Source,
	}), &gorm.Config{
		Logger: NewGormLogger(),
	})
	if err != nil {
彭芳 committed
45
		return nil, err
彭芳 committed
46
	}
彭芳 committed
47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63
	sqlDb, err := db.DB()
	if err != nil {
		return nil, err
	}
	// the maximum number of connections in the idle connection pool.
	if c.Database.MaxIdle > 0 {
		sqlDb.SetMaxIdleConns(int(c.Database.MaxIdle))
	}
	// the maximum number of open connections to the database.
	if c.Database.MaxOpen > 0 {
		sqlDb.SetMaxOpenConns(int(c.Database.MaxOpen))
	}
	// the maximum amount of time a connection may be reused.
	if c.Database.MaxLifetime != nil {
		sqlDb.SetConnMaxLifetime(c.Database.MaxLifetime.AsDuration())
	}
	return db, nil
彭芳 committed
64 65 66 67 68 69
}

func NewGormLogger() logger.Interface {
	return logger.New(logrus.StandardLogger(), logger.Config{
		SlowThreshold:             time.Second,
		IgnoreRecordNotFoundError: true,
彭芳 committed
70
		LogLevel:                  logger.Info,
彭芳 committed
71 72
	})
}