308 lines
7.8 KiB
Go
308 lines
7.8 KiB
Go
package db
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"fmt"
|
|
"os"
|
|
"runtime/debug"
|
|
"strings"
|
|
"time"
|
|
|
|
"91porn-server/common/log"
|
|
|
|
"go.mongodb.org/mongo-driver/mongo"
|
|
"go.mongodb.org/mongo-driver/mongo/options"
|
|
"go.mongodb.org/mongo-driver/x/mongo/driver/connstring"
|
|
)
|
|
|
|
var dataSource = make(map[string]*MongoDB)
|
|
var registerPool []Register
|
|
|
|
// DBmap
|
|
type DBmap struct {
|
|
Key string
|
|
URL string
|
|
}
|
|
|
|
type Register struct {
|
|
Key string
|
|
Table []string
|
|
}
|
|
|
|
// MongoOptions 数据库配置
|
|
type MongoOptions struct {
|
|
URL string `json:"url"` // 服务器连接地址
|
|
}
|
|
|
|
// MongoDB MongoDB
|
|
type MongoDB struct {
|
|
db *mongo.Database
|
|
}
|
|
|
|
func (m *MongoDB) Tool() *MongoTool {
|
|
return m.ToolCtx(context.Background())
|
|
}
|
|
|
|
func (m *MongoDB) ToolCtx(ctx context.Context) *MongoTool {
|
|
return &MongoTool{db: m.db, ctx: ctx}
|
|
}
|
|
|
|
// Coll 获取表名
|
|
func (m *MongoDB) Coll(name string) *MongoTool {
|
|
t := m.Tool()
|
|
return t.Coll(name)
|
|
}
|
|
|
|
// CollCtx 获取表名 从外部传入ctx
|
|
func (m *MongoDB) CollCtx(ctx context.Context, name string) *MongoTool {
|
|
t := m.ToolCtx(ctx)
|
|
return t.Coll(name)
|
|
}
|
|
|
|
// MongoTool mongo官方库事务封装
|
|
type MongoTool struct {
|
|
db *mongo.Database
|
|
ctx context.Context // 当前使用的ctx
|
|
coll *mongo.Collection
|
|
}
|
|
|
|
func (m *MongoTool) Coll(name string) *MongoTool {
|
|
opt := options.CollectionOptions{Registry: registry}
|
|
m.coll = m.db.Collection(name, &opt)
|
|
return m
|
|
}
|
|
|
|
// Trans 开启事务处理包裹处理,里面处理的全是利用的事务的ctx
|
|
func (m *MongoDB) Trans(fn func(*MongoTool) error, opts ...*TransOpts) error {
|
|
return m.TransCtx(context.Background(), fn, opts...)
|
|
}
|
|
|
|
// TransCtx 外部传入ctx
|
|
func (m *MongoDB) TransCtx(ctx context.Context, fn func(*MongoTool) error, opts ...*TransOpts) error {
|
|
t := m.ToolCtx(ctx)
|
|
return m.db.Client().UseSession(ctx, func(sessionContext mongo.SessionContext) error {
|
|
if err := sessionContext.StartTransaction(); err != nil {
|
|
return err
|
|
}
|
|
defer func() {
|
|
if err := recover(); err != nil {
|
|
log.Warn(fmt.Sprintf("caught panic during transaction, aborting. err: %+v, stack: %s", err, debug.Stack()))
|
|
if err := sessionContext.AbortTransaction(sessionContext); err != nil {
|
|
log.Warn("mongo AbortTransaction panic err", log.E(err))
|
|
}
|
|
}
|
|
sessionContext.EndSession(sessionContext)
|
|
}()
|
|
t.ctx = sessionContext
|
|
if err := runTransactionWithRetry(t, fn, MergeTransOpts(opts)); err != nil {
|
|
if strings.Contains(err.Error(), "NoSuchTransaction") {
|
|
log.Warn("NoSuchTransaction error, return")
|
|
return err
|
|
}
|
|
if err := sessionContext.AbortTransaction(sessionContext); err != nil {
|
|
log.Warn("mongo AbortTransaction err", log.E(err))
|
|
}
|
|
log.Warn("caught exception during transaction, aborting.", log.E(err))
|
|
sessionContext.EndSession(sessionContext)
|
|
return err
|
|
}
|
|
return commitWithRetry(sessionContext)
|
|
})
|
|
}
|
|
|
|
// runTransactionWithRetry is an example function demonstrating transaction retry logic.
|
|
func runTransactionWithRetry(t *MongoTool, txnFn func(t *MongoTool) error, opts *TransOpts) error {
|
|
//no set ReEntryCount is loop retry
|
|
for {
|
|
err := txnFn(t) // Performs transaction.
|
|
if err == nil {
|
|
return nil
|
|
}
|
|
log.Warn("Transaction aborted. Caught exception during transaction.", log.E(err))
|
|
// If transient error, retry the whole transaction
|
|
if strings.Contains(err.Error(), "NoSuchTransaction") {
|
|
log.Info("NoSuchTransaction error,return and break retry loop")
|
|
return err
|
|
}
|
|
cmdErr, ok := err.(mongo.CommandError)
|
|
if ok && cmdErr.HasErrorLabel("TransientTransactionError") {
|
|
if opts != nil {
|
|
if opts.ReEntryCount != nil {
|
|
if *opts.ReEntryCount <= 0 {
|
|
return cmdErr
|
|
}
|
|
*opts.ReEntryCount--
|
|
}
|
|
}
|
|
log.Info("TransientTransactionError, retrying transaction...")
|
|
continue
|
|
}
|
|
// else return err
|
|
return err
|
|
}
|
|
}
|
|
|
|
// commitWithRetry is an example function demonstrating transaction retry logic.
|
|
func commitWithRetry(sess mongo.SessionContext) error {
|
|
for {
|
|
err := sess.CommitTransaction(sess)
|
|
switch e := err.(type) {
|
|
case nil:
|
|
log.Info("Transaction committed.")
|
|
return nil
|
|
case mongo.CommandError:
|
|
// Can retry commit
|
|
if e.HasErrorLabel("UnknownTransactionCommitResult") {
|
|
log.Info("UnknownTransactionCommitResult, retrying commit operation...")
|
|
continue
|
|
}
|
|
log.Info("Error during commit...")
|
|
return e
|
|
default:
|
|
log.Info("Error during commit...")
|
|
return e
|
|
}
|
|
}
|
|
}
|
|
|
|
// Disconnect 关闭数据库连接
|
|
func (m *MongoDB) Disconnect() error {
|
|
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
|
log.Info("closing mongodb connections")
|
|
defer cancel()
|
|
if err := m.db.Client().Disconnect(ctx); err != nil {
|
|
log.Warn(fmt.Sprintf("close mongo connections err: %+v", err))
|
|
return err
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// InitMongo 初始化Mongo
|
|
func InitMongo(murl string) (*MongoDB, error) {
|
|
cs, err := connstring.Parse(murl)
|
|
if err != nil {
|
|
log.Error("mongo URL parse fail", log.Any("url", murl), log.E(err))
|
|
return nil, err
|
|
}
|
|
db := cs.Database
|
|
ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second)
|
|
defer cancel()
|
|
// Connect to MongoDB
|
|
mongoClient, err := mongo.Connect(ctx, options.Client().ApplyURI(murl))
|
|
if err != nil {
|
|
log.Error("mongodb connect fail", log.E(err))
|
|
return nil, err
|
|
}
|
|
// Check the connection
|
|
if err = mongoClient.Ping(context.Background(), nil); err != nil {
|
|
log.Error("mongodb connect ping is fail")
|
|
return nil, err
|
|
}
|
|
mongoDataBase := mongoClient.Database(db)
|
|
return &MongoDB{db: mongoDataBase}, nil
|
|
}
|
|
|
|
func InitDS(dbmap []DBmap, register []Register) map[string]*MongoDB {
|
|
if len(dbmap) == 0 {
|
|
log.Error("dbmap must be not empty ")
|
|
os.Exit(1)
|
|
}
|
|
for _, v := range dbmap {
|
|
db, err := InitMongo(v.URL)
|
|
if err != nil {
|
|
log.Error(fmt.Sprintf("[DB-%s] start up error", v.Key), log.E(err))
|
|
os.Exit(1)
|
|
}
|
|
log.Info(fmt.Sprintf("[DB-%s] connect is successfully", v.Key))
|
|
dataSource[v.Key] = db
|
|
}
|
|
registerPool = register
|
|
return dataSource
|
|
}
|
|
|
|
func Init(tableName string) *MongoDB {
|
|
if tableName == "" {
|
|
panic(errors.New("table name must not be empty"))
|
|
}
|
|
return Selector(tableName)
|
|
}
|
|
|
|
// 初始化只读数据
|
|
func InitRead(tableName string) *MongoDB {
|
|
if tableName == "" {
|
|
panic(errors.New("table name must not be empty"))
|
|
}
|
|
return SelectorRead(tableName)
|
|
}
|
|
|
|
func CloseDS() {
|
|
if len(dataSource) == 0 {
|
|
log.Warn("dataSource empty ")
|
|
return
|
|
}
|
|
for k, v := range dataSource {
|
|
if err := v.Disconnect(); err != nil {
|
|
log.Error(fmt.Sprintf("[DB-%s]Mongo Disconnect error", k), log.E(err))
|
|
continue
|
|
}
|
|
log.Info(fmt.Sprintf("[DB-%s]Mongo Disconnect OK", k))
|
|
}
|
|
}
|
|
|
|
// BaseDAO 如果ctx为nil 则表示不使用事务,如果ctx不为空则表示使用事务
|
|
func BaseDAO(tableName string, ctx context.Context) *MongoTool {
|
|
if tableName == "" {
|
|
panic(errors.New("table name must not be empty"))
|
|
}
|
|
mongdb := Selector(tableName)
|
|
if ctx == nil {
|
|
return &MongoTool{db: mongdb.db, ctx: context.Background(), coll: mongdb.db.Collection(tableName)}
|
|
}
|
|
return &MongoTool{db: mongdb.db, ctx: ctx, coll: mongdb.db.Collection(tableName)}
|
|
}
|
|
|
|
func Selector(tableName string) *MongoDB {
|
|
var isExist = false
|
|
var db *MongoDB
|
|
for _, r := range registerPool {
|
|
if strings.HasPrefix(r.Key, "Read") {
|
|
continue
|
|
}
|
|
for _, v := range r.Table {
|
|
if v == tableName {
|
|
isExist = true
|
|
db = dataSource[r.Key]
|
|
break
|
|
}
|
|
}
|
|
}
|
|
if !isExist {
|
|
log.Warn("current table not register,please register it first", log.Any("table", tableName))
|
|
panic(errors.New("current table not register"))
|
|
}
|
|
return db
|
|
}
|
|
|
|
func SelectorRead(tableName string) *MongoDB {
|
|
var isExist = false
|
|
var db *MongoDB
|
|
for _, r := range registerPool {
|
|
if !strings.HasPrefix(r.Key, "Read") {
|
|
continue
|
|
}
|
|
for _, v := range r.Table {
|
|
if v == tableName {
|
|
isExist = true
|
|
db = dataSource[r.Key]
|
|
break
|
|
}
|
|
}
|
|
}
|
|
if !isExist {
|
|
log.Warn("current table not register,please register it first", log.Any("table", tableName))
|
|
panic(errors.New("current table not register"))
|
|
}
|
|
return db
|
|
}
|