ku-kratos-layout/internal/data/data.go

package data

import (
	"github.com/go-kratos/kratos-layout/internal/conf"
	"github.com/go-kratos/kratos-layout/internal/data/models"
	"github.com/go-kratos/kratos/v2/log"
	"github.com/google/wire"
	"github.com/redis/go-redis/v9"
)

// ProviderSet is data providers.
var ProviderSet = wire.NewSet(NewRedis, NewData)

// Data .
type Data struct {
	log *log.Helper
	rdb *redis.Client
	*models.Mysql
}

// NewData .
func NewData(c *conf.Bootstrap, rdb *redis.Client, logger log.Logger) (*Data, func(), error) {
	cleanup := func() {
	}
	data := &Data{
		log: log.NewHelper(log.With(logger, "module", "data/data")),
	}
	var err error
	if c.Data.Database != nil {
		data.Mysql, err = models.NewMysql(c.Data.Database.Source, logger)
		if err != nil {
			return nil, cleanup, err
		}
		data.Mysql.XDB.SetMaxIdleConns(int(c.Data.Database.MaxIdleConns))
		data.Mysql.XDB.SetMaxOpenConns(int(c.Data.Database.MaxOpenConns))
	}
	data.rdb = rdb
	cleanup = func() {
		if data.XDB != nil {
			data.XDB.Close()
		}
		if data.rdb != nil {
			data.rdb.Close()
		}
	}
	return data, cleanup, nil
}

func NewRedis(c *conf.Bootstrap) *redis.Client {
	if c.Data.Redis == nil {
		return nil
	}
	rdb := redis.NewClient(&redis.Options{
		Addr:         c.Data.Redis.Addr,
		DB:           int(c.Data.Redis.Db),
		DialTimeout:  c.Data.Redis.DialTimeout.AsDuration(),
		WriteTimeout: c.Data.Redis.WriteTimeout.AsDuration(),
		ReadTimeout:  c.Data.Redis.ReadTimeout.AsDuration(),
		PoolSize:     int(c.Data.Redis.PoolSize),
		MinIdleConns: int(c.Data.Redis.MinIdleConns),
	})
	return rdb
}