Claude Opus 5
@@ -0,0 +1,307 @@
|
||||
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
|
||||
}
|
||||
Reference in New Issue
Block a user