100 lines
1.7 KiB
Go
100 lines
1.7 KiB
Go
//
|
|
// mongo.go
|
|
// Copyright (C) 2023 tiglog <me@tiglog.com>
|
|
//
|
|
// Distributed under terms of the MIT license.
|
|
//
|
|
|
|
package mgodb
|
|
|
|
import (
|
|
"context"
|
|
"time"
|
|
|
|
"go.mongodb.org/mongo-driver/mongo"
|
|
"go.mongodb.org/mongo-driver/mongo/options"
|
|
)
|
|
|
|
var Mm *sMongoManager
|
|
|
|
type sMongoManager struct {
|
|
cli *mongo.Client
|
|
name string
|
|
poolSize int
|
|
uri string
|
|
}
|
|
|
|
func Init(uri, dbName string, poolSize int) {
|
|
Mm = &sMongoManager{
|
|
name: dbName,
|
|
uri: uri,
|
|
poolSize: poolSize,
|
|
}
|
|
}
|
|
|
|
func (s *sMongoManager) SetDb(name string) {
|
|
s.name = name
|
|
}
|
|
|
|
func (s *sMongoManager) Connect() error {
|
|
var err error
|
|
clientOptions := options.Client().ApplyURI(s.uri)
|
|
clientOptions.SetMaxPoolSize(uint64(s.poolSize))
|
|
|
|
// 连接到MongoDB
|
|
s.cli, err = mongo.Connect(context.TODO(), clientOptions)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
// 检查连接
|
|
err = s.cli.Ping(context.TODO(), nil)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
return nil
|
|
}
|
|
func (s *sMongoManager) Client() (*mongo.Client, error) {
|
|
if s.cli == nil {
|
|
err := s.Connect()
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
}
|
|
return s.cli, nil
|
|
}
|
|
|
|
func (s *sMongoManager) Db() *mongo.Database {
|
|
c, err := s.Client()
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
return c.Database(s.name)
|
|
}
|
|
|
|
func (s *sMongoManager) Collection(name string) *mongo.Collection {
|
|
c, err := s.Client()
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
return c.Database(s.name).Collection(name)
|
|
}
|
|
|
|
func NewCtx() context.Context {
|
|
return context.Background()
|
|
}
|
|
|
|
func NewTtlCtx(ttl time.Duration) (context.Context, context.CancelFunc) {
|
|
return context.WithTimeout(context.Background(), ttl)
|
|
}
|
|
|
|
func (s *sMongoManager) Close() error {
|
|
if s.cli == nil {
|
|
return nil
|
|
}
|
|
err := s.cli.Disconnect(NewCtx())
|
|
if err != nil {
|
|
return err
|
|
}
|
|
return nil
|
|
}
|