wxbot-layout/internal/data/data.go

package data

import (
	"context"
	"fmt"
	"time"

	"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"
	"go.mongodb.org/mongo-driver/mongo"
	"go.mongodb.org/mongo-driver/mongo/options"
	"go.mongodb.org/mongo-driver/mongo/readconcern"
	"go.mongodb.org/mongo-driver/mongo/writeconcern"
)

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

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

// NewData .
func NewData(c *conf.Bootstrap, rdb *redis.Client, mongodb *mongo.Database, logger log.Logger, ctx context.Context) (*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.SetConnMaxIdleTime(time.Minute * 10)
		data.Mysql.XDB.SetMaxOpenConns(int(c.Data.Database.MaxOpenConns))
	}
	if rdb != nil {
		data.rdb = rdb
	}
	if mongodb != nil {
		data.mongodb = mongodb
	}
	cleanup = func() {
		if data.XDB != nil {
			data.XDB.Close()
		}
		if data.rdb != nil {
			data.rdb.Close()
		}
		if data.mongodb != nil {
			data.mongodb.Client().Disconnect(ctx)
		}
	}
	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),
	})
	key := fmt.Sprintf("%s:%s", c.Server.Registrar.Group, c.Server.Registrar.Data)
	rdb.Set(context.Background(), fmt.Sprintf("used_by_%s", key), 1, -1)
	return rdb
}

func NewMongo(ctx context.Context, c *conf.Bootstrap) (*mongo.Database, error) {
	if c.Data.Mongo == nil {
		return nil, nil
	}
	clientOptions := options.Client().ApplyURI(c.Data.Mongo.Uri).SetSocketTimeout(time.Second * 30).
		SetReadConcern(readconcern.Majority()).SetWriteConcern(writeconcern.New(writeconcern.WMajority(), writeconcern.J(true)))
	client, err := mongo.Connect(ctx, clientOptions)

	if err != nil {
		return nil, err
	}
	database := client.Database(c.Data.Mongo.Database)
	return database, nil
}