forked from all_user/stat_center
139 lines
3.1 KiB
Go
139 lines
3.1 KiB
Go
package mongos
|
|
|
|
import (
|
|
"context"
|
|
"stat/common/log"
|
|
"time"
|
|
|
|
"go.mongodb.org/mongo-driver/bson"
|
|
"go.mongodb.org/mongo-driver/x/mongo/driver/connstring"
|
|
|
|
"go.mongodb.org/mongo-driver/mongo"
|
|
"go.mongodb.org/mongo-driver/mongo/options"
|
|
)
|
|
|
|
var MgoManager *Manager
|
|
|
|
type Options struct {
|
|
Url string `json:"url"` // 数据库连接
|
|
}
|
|
|
|
// Manager 数据库连接管理器
|
|
// 重置连接等
|
|
type Manager struct {
|
|
client *mongo.Client
|
|
url string
|
|
db string
|
|
}
|
|
|
|
func NewManager(url string) *Manager {
|
|
return &Manager{url: url}
|
|
}
|
|
|
|
func (m *Manager) Database() *mongo.Database {
|
|
return m.client.Database(m.db)
|
|
}
|
|
|
|
func (m *Manager) Connect() error {
|
|
cs, err := connstring.Parse(m.url)
|
|
if err != nil {
|
|
log.Error("db connect parse err", log.E(err))
|
|
return err
|
|
}
|
|
|
|
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
|
defer cancel()
|
|
|
|
client, err := mongo.Connect(ctx, options.Client().ApplyURI(m.url))
|
|
if err != nil {
|
|
log.Error("[warn] connect mongo db err", log.E(err))
|
|
return err
|
|
}
|
|
if err = client.Ping(ctx, nil); err != nil {
|
|
log.Error("[warn] ping mongo db err", log.E(err))
|
|
return err
|
|
}
|
|
|
|
m.db = cs.Database
|
|
m.client = client
|
|
|
|
log.Info("mongos:Init ok!", log.Any("dbName", m.db), log.Any("dbHost", cs.Hosts))
|
|
return nil
|
|
}
|
|
|
|
func (m *Manager) DisConnect() {
|
|
if m.client != nil {
|
|
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
|
defer cancel()
|
|
if err := m.client.Disconnect(ctx); err != nil {
|
|
log.Error("db disconnect err", log.E(err))
|
|
}
|
|
m.client = nil
|
|
}
|
|
}
|
|
|
|
//Transaction 开启事务处理包裹处理,里面处理的全是利用的事务的ctx
|
|
func (m *Manager) Transaction(ctx context.Context, fn func(sctx mongo.SessionContext) error) error {
|
|
return m.client.UseSession(
|
|
ctx,
|
|
func(sctx mongo.SessionContext) error {
|
|
return runTransactionWithRetry(sctx, func(ctx mongo.SessionContext) error {
|
|
err := sctx.StartTransaction()
|
|
if err != nil {
|
|
return err
|
|
}
|
|
err = fn(sctx)
|
|
if err != nil {
|
|
sctx.AbortTransaction(sctx)
|
|
return err
|
|
}
|
|
return commitWithRetry(sctx)
|
|
})
|
|
},
|
|
)
|
|
}
|
|
|
|
func runTransactionWithRetry(sctx mongo.SessionContext, txnFn func(mongo.SessionContext) error) error {
|
|
for {
|
|
err := txnFn(sctx) // Performs transaction.
|
|
if err == nil {
|
|
return nil
|
|
}
|
|
// If transient error, retry the whole transaction
|
|
if cmdErr, ok := err.(mongo.CommandError); ok && cmdErr.HasErrorLabel("TransientTransactionError") {
|
|
continue
|
|
}
|
|
return err
|
|
}
|
|
}
|
|
|
|
// commitWithRetry is an example function demonstrating transaction retry logic.
|
|
func commitWithRetry(sctx mongo.SessionContext) error {
|
|
for {
|
|
err := sctx.CommitTransaction(sctx)
|
|
switch e := err.(type) {
|
|
case nil:
|
|
return nil
|
|
case mongo.CommandError:
|
|
// Can retry commit
|
|
if e.HasErrorLabel("UnknownTransactionCommitResult") {
|
|
continue
|
|
}
|
|
return e
|
|
default:
|
|
return e
|
|
}
|
|
}
|
|
}
|
|
|
|
func Transaction(ctx context.Context, fn func(sctx mongo.SessionContext) error) error {
|
|
return MgoManager.Transaction(ctx, fn)
|
|
}
|
|
|
|
func ToBsonM(v interface{}) bson.M {
|
|
b, _ := bson.Marshal(v)
|
|
var r = bson.M{}
|
|
_ = bson.Unmarshal(b, &r)
|
|
return r
|
|
}
|