49 lines
1.2 KiB
Go
49 lines
1.2 KiB
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
|
|
c *config.Config
|
|
}
|
|
|
|
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))
|
|
}
|
|
|
|
if err = transDBClient.Ping(ctx, nil); err != nil {
|
|
panic(fmt.Sprintf("mongo链接失败: %v", err))
|
|
}
|
|
return &Mongo{
|
|
Client: transDBClient,
|
|
c: c,
|
|
}, func() {
|
|
transDBClient.Disconnect(ctx)
|
|
}
|
|
}
|
|
|
|
type MongoModel interface {
|
|
MongoTableName() string
|
|
}
|
|
|
|
func (m *Mongo) Co(mongoModel MongoModel) *mongo.Collection {
|
|
return m.Client.Database(m.c.Mongo.DataBase).Collection(mongoModel.MongoTableName())
|
|
}
|