package data import ( "context" "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/go-redis/redis/v8" "github.com/google/wire" "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.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), }) 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 }