ai_scheduler/internal/pkg/mongo.go

36 lines
868 B
Go

package pkg
import (
"ai_scheduler/internal/config"
"ai_scheduler/internal/pkg/utils_mongo"
"context"
"fmt"
"time"
"go.mongodb.org/mongo-driver/mongo"
)
type Mongo struct {
Client *mongo.Client
}
func NewMongoDb(ctx context.Context, c *config.Config) (*Mongo, func()) {
transDBClient, err := utils_mongo.NewMongoClient(ctx, &utils_mongo.ClientStruct{
Uri: c.Mongo.Source,
MaxPoolSize: c.Mongo.MaxPoolSize,
MinPoolSize: c.Mongo.MinPoolSize,
MaxConnIdleTime: time.Duration(c.Mongo.MaxConnIdleTime) * time.Minute,
ConnectTimeout: time.Duration(c.Mongo.ConnectTimeout) * time.Second,
SocketTimeout: time.Duration(c.Mongo.SocketTimeout) * time.Second,
})
if err != nil {
panic(fmt.Sprintf("mongo数据库错误: %v", err))
}
return &Mongo{
Client: transDBClient,
}, func() {
transDBClient.Disconnect(ctx)
}
}