Initial commit

common/mongos/manager.go

@@ -0,0 +1,138 @@
+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
+}