dorm/mongo_session.go

package dorm

import (
	"context"
	"errors"
	"fmt"
	"go.mongodb.org/mongo-driver/mongo"
)

type MongoSessionOptions struct {
	db                 *mongo.Client        // 驱动
	configDatabaseName string               // 库名
	session            mongo.Session        // 会话
	sessionContext     mongo.SessionContext // 会话上下文
}

// Begin 开始事务，会同时创建开始会话需要在退出时关闭会话
func (c *MongoClient) Begin(ctx context.Context) (ms *MongoSessionOptions, err error) {

	ms.db = c.db
	ms.configDatabaseName = c.configDatabaseName

	// 开始会话
	ms.session, err = ms.db.StartSession()
	if err != nil {
		return nil, errors.New(fmt.Sprintf("开始会话失败：%v", err))
	}

	// 会话上下文
	ms.sessionContext = mongo.NewSessionContext(ctx, ms.session)

	// 会话开启事务
	err = ms.session.StartTransaction()

	return ms, err
}

// Close 关闭会话
func (cs *MongoSessionOptions) Close(ctx context.Context) {
	cs.session.EndSession(ctx)
}

// Rollback 回滚事务
func (cs *MongoSessionOptions) Rollback(ctx context.Context) error {
	return cs.session.AbortTransaction(ctx)
}

// Commit 提交事务
func (cs *MongoSessionOptions) Commit(ctx context.Context) error {
	return cs.session.CommitTransaction(ctx)
}