Claude Opus 4.7
@@ -0,0 +1,150 @@
|
||||
package mongos
|
||||
|
||||
import (
|
||||
"go.mongodb.org/mongo-driver/bson"
|
||||
)
|
||||
|
||||
type QueryBuilder struct {
|
||||
query bson.M
|
||||
}
|
||||
|
||||
func NewQueryBuilder() *QueryBuilder {
|
||||
return &QueryBuilder{query: bson.M{}}
|
||||
}
|
||||
|
||||
func (q *QueryBuilder) Build() bson.M {
|
||||
return q.query
|
||||
}
|
||||
func (q *QueryBuilder) addCondition(condition bool, field string, value interface{}) {
|
||||
if condition {
|
||||
q.query[field] = value
|
||||
}
|
||||
}
|
||||
|
||||
func (q *QueryBuilder) Eq(field string, value interface{}, condition ...bool) *QueryBuilder {
|
||||
if len(condition) == 0 || condition[0] {
|
||||
q.query[field] = value
|
||||
}
|
||||
return q
|
||||
}
|
||||
|
||||
func (q *QueryBuilder) Ne(field string, value interface{}, condition ...bool) *QueryBuilder {
|
||||
if len(condition) == 0 || condition[0] {
|
||||
q.query[field] = bson.M{"$ne": value}
|
||||
}
|
||||
return q
|
||||
}
|
||||
|
||||
func (q *QueryBuilder) Or(conditions ...bson.M) *QueryBuilder {
|
||||
if len(conditions) > 0 {
|
||||
q.query["$or"] = conditions
|
||||
}
|
||||
return q
|
||||
}
|
||||
|
||||
func (q *QueryBuilder) Set(value bson.M) *QueryBuilder {
|
||||
q.query["$set"] = value
|
||||
return q
|
||||
}
|
||||
|
||||
func (q *QueryBuilder) Gt(field string, value interface{}, condition ...bool) *QueryBuilder {
|
||||
if len(condition) == 0 || condition[0] {
|
||||
q.query[field] = bson.M{"$gt": value}
|
||||
}
|
||||
return q
|
||||
}
|
||||
|
||||
func (q *QueryBuilder) Gte(field string, value interface{}, condition ...bool) *QueryBuilder {
|
||||
if len(condition) == 0 || condition[0] {
|
||||
q.query[field] = bson.M{"$gte": value}
|
||||
}
|
||||
return q
|
||||
}
|
||||
|
||||
func (q *QueryBuilder) Lt(field string, value interface{}, condition ...bool) *QueryBuilder {
|
||||
if len(condition) == 0 || condition[0] {
|
||||
q.query[field] = bson.M{"$lt": value}
|
||||
}
|
||||
return q
|
||||
}
|
||||
|
||||
func (q *QueryBuilder) Lte(field string, value interface{}, condition ...bool) *QueryBuilder {
|
||||
if len(condition) == 0 || condition[0] {
|
||||
q.query[field] = bson.M{"$lte": value}
|
||||
}
|
||||
return q
|
||||
}
|
||||
|
||||
func (q *QueryBuilder) In(field string, values *[]int64, condition ...bool) *QueryBuilder {
|
||||
if len(condition) == 0 || condition[0] {
|
||||
q.query[field] = bson.M{"$in": *values}
|
||||
}
|
||||
return q
|
||||
}
|
||||
|
||||
func (q *QueryBuilder) Nin(field string, values []interface{}, condition ...bool) *QueryBuilder {
|
||||
if len(condition) == 0 || condition[0] {
|
||||
q.query[field] = bson.M{"$nin": values}
|
||||
}
|
||||
return q
|
||||
}
|
||||
|
||||
func (q *QueryBuilder) Regex(field string, value interface{}, condition ...bool) *QueryBuilder {
|
||||
if len(condition) == 0 || condition[0] {
|
||||
q.query[field] = bson.M{"regex": value.(string)}
|
||||
}
|
||||
return q
|
||||
}
|
||||
|
||||
func (q *QueryBuilder) All(field string, values []interface{}, condition ...bool) *QueryBuilder {
|
||||
if len(condition) == 0 || condition[0] {
|
||||
q.query[field] = bson.M{"$all": values}
|
||||
}
|
||||
return q
|
||||
}
|
||||
|
||||
func (q *QueryBuilder) ElemMatch(field string, value interface{}, condition ...bool) *QueryBuilder {
|
||||
if len(condition) == 0 || condition[0] {
|
||||
q.query[field] = bson.M{"$elemMatch": value}
|
||||
}
|
||||
return q
|
||||
}
|
||||
|
||||
func (q *QueryBuilder) Size(field string, value int, condition ...bool) *QueryBuilder {
|
||||
if len(condition) == 0 || condition[0] {
|
||||
q.query[field] = bson.M{"$size": value}
|
||||
}
|
||||
return q
|
||||
}
|
||||
|
||||
func (q *QueryBuilder) Exists(field string, value bool, condition ...bool) *QueryBuilder {
|
||||
if len(condition) == 0 || condition[0] {
|
||||
q.query[field] = bson.M{"$exists": value}
|
||||
}
|
||||
return q
|
||||
}
|
||||
|
||||
func (q *QueryBuilder) Mod(field string, divisor, remainder int, condition ...bool) *QueryBuilder {
|
||||
if len(condition) == 0 || condition[0] {
|
||||
q.query[field] = bson.M{"$mod": []int{divisor, remainder}}
|
||||
}
|
||||
return q
|
||||
}
|
||||
|
||||
func (q *QueryBuilder) GeoWithin(field string, geometry interface{}, condition ...bool) *QueryBuilder {
|
||||
if len(condition) == 0 || condition[0] {
|
||||
q.query[field] = bson.M{"$geoWithin": bson.M{
|
||||
"$geometry": geometry,
|
||||
}}
|
||||
}
|
||||
return q
|
||||
}
|
||||
|
||||
func (q *QueryBuilder) GeoIntersects(field string, geometry interface{}, condition ...bool) *QueryBuilder {
|
||||
if len(condition) == 0 || condition[0] {
|
||||
q.query[field] = bson.M{"$geoIntersects": bson.M{
|
||||
"$geometry": geometry,
|
||||
}}
|
||||
}
|
||||
return q
|
||||
}
|
||||
@@ -0,0 +1,582 @@
|
||||
package mongos
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"videoweb/pub/common/log"
|
||||
"videoweb/pub/common/utils"
|
||||
|
||||
"reflect"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"go.mongodb.org/mongo-driver/bson"
|
||||
"go.mongodb.org/mongo-driver/bson/primitive"
|
||||
"go.mongodb.org/mongo-driver/mongo"
|
||||
"go.mongodb.org/mongo-driver/mongo/options"
|
||||
)
|
||||
|
||||
type Collection struct {
|
||||
name string
|
||||
coll *mongo.Collection
|
||||
}
|
||||
|
||||
func (b *Collection) FindOne(ctx context.Context, bind interface{}, query bson.M, opts ...*options.FindOneOptions) error {
|
||||
err := b.coll.FindOne(ctx, query, opts...).Decode(bind)
|
||||
if err != nil && err == mongo.ErrNoDocuments {
|
||||
return nil
|
||||
}
|
||||
return err
|
||||
}
|
||||
|
||||
func (b *Collection) FindOneEx(ctx context.Context, bind interface{}, query bson.M, opts ...*options.FindOneOptions) error {
|
||||
err := b.coll.FindOne(ctx, query, opts...).Decode(bind)
|
||||
return err
|
||||
}
|
||||
|
||||
func (b *Collection) FindMOneAndUpdate(ctx context.Context, bind interface{}, query bson.M, update interface{}, opts ...*options.FindOneAndUpdateOptions) error {
|
||||
setReturn := false
|
||||
for _, opt := range opts {
|
||||
if opt != nil && opt.ReturnDocument != nil {
|
||||
setReturn = true
|
||||
break
|
||||
}
|
||||
}
|
||||
if !setReturn {
|
||||
after := options.After
|
||||
opts = append(opts, &options.FindOneAndUpdateOptions{ReturnDocument: &after})
|
||||
}
|
||||
um := utils.Struct2Map(update)
|
||||
// um["update_at"] = time.Now()
|
||||
// set := bson.M{
|
||||
// "$set": um,
|
||||
// }
|
||||
result := b.coll.FindOneAndUpdate(ctx, query, um, opts...)
|
||||
err := result.Decode(bind)
|
||||
return err
|
||||
}
|
||||
|
||||
func (b *Collection) FindMOneAndUpdateInc(ctx context.Context, bind interface{}, query bson.M, update interface{}, inc interface{}, opts ...*options.FindOneAndUpdateOptions) error {
|
||||
um := utils.Struct2Map(update)
|
||||
um["update_at"] = time.Now()
|
||||
im := utils.Struct2Map(inc)
|
||||
set := bson.M{
|
||||
"$set": um,
|
||||
}
|
||||
if im != nil {
|
||||
set["$inc"] = inc
|
||||
}
|
||||
|
||||
setReturn := false
|
||||
for _, opt := range opts {
|
||||
if opt != nil && opt.ReturnDocument != nil {
|
||||
setReturn = true
|
||||
break
|
||||
}
|
||||
}
|
||||
if !setReturn {
|
||||
after := options.After
|
||||
opts = append(opts, &options.FindOneAndUpdateOptions{ReturnDocument: &after})
|
||||
}
|
||||
|
||||
result := b.coll.FindOneAndUpdate(ctx, query, set, opts...)
|
||||
err := result.Decode(bind)
|
||||
return err
|
||||
}
|
||||
|
||||
// func (b *Collection) FindOneAndUpdateTiny(
|
||||
// ctx context.Context,
|
||||
// query interface{},
|
||||
// update bson.M,
|
||||
// opts ...*options.FindOneAndUpdateOptions) error {
|
||||
// result := b.collection().FindOneAndUpdate(ctx, query, update, opts...)
|
||||
// return result.Err()
|
||||
// }
|
||||
|
||||
// func (b *MongoTool) FindOneAndUpdateReturnTiny(ctx context.Context,
|
||||
// bind interface{},
|
||||
// query bson.M,
|
||||
// update bson.M,
|
||||
// afterDoc bool,
|
||||
// opts ...*options.FindOneAndUpdateOptions) error {
|
||||
|
||||
// if afterDoc {
|
||||
// opts = append(opts, options.FindOneAndUpdate().SetReturnDocument(options.After))
|
||||
// } else {
|
||||
// opts = append(opts, options.FindOneAndUpdate().SetReturnDocument(options.Before))
|
||||
// }
|
||||
|
||||
// result := b.collection().FindOneAndUpdate(ctx, query, update, opts...).Decode(bind)
|
||||
// return result
|
||||
// }
|
||||
|
||||
// func (b *Collection) FindOneAndUpdate(ctx context.Context, bind interface{}, query bson.M, update bson.M, opts ...*options.FindOneAndUpdateOptions) error {
|
||||
// setReturn := false
|
||||
// for _, opt := range opts {
|
||||
// if opt != nil && opt.ReturnDocument != nil {
|
||||
// setReturn = true
|
||||
// break
|
||||
// }
|
||||
// }
|
||||
// if !setReturn {
|
||||
// after := options.After
|
||||
// opts = append(opts, &options.FindOneAndUpdateOptions{ReturnDocument: &after})
|
||||
// }
|
||||
// b.setUpdatedAt(update, time.Now())
|
||||
// result := b.collection().FindOneAndUpdate(ctx, query, update, opts...)
|
||||
// err := result.Decode(bind)
|
||||
// return err
|
||||
// }
|
||||
|
||||
func (b *Collection) Find(ctx context.Context, bind interface{}, query bson.M, opt ...*options.FindOptions) error {
|
||||
err := b.validSlice(bind)
|
||||
cur, err := b.coll.Find(ctx, query, opt...)
|
||||
if err != nil && err == mongo.ErrNoDocuments {
|
||||
return nil
|
||||
} else if err != nil {
|
||||
return err
|
||||
}
|
||||
err = cur.All(ctx, bind)
|
||||
return err
|
||||
}
|
||||
|
||||
// func (b *Collection) FindByPage(ctx context.Context, bind interface{}, query bson.M, pageAndSort *PageAndSort, opts ...*options.FindOptions) error {
|
||||
// err := b.validSlice(bind)
|
||||
|
||||
// if err != nil {
|
||||
// return err
|
||||
// }
|
||||
|
||||
// opt := &options.FindOptions{}
|
||||
// if pageAndSort == nil {
|
||||
// return errors.New("page must not nil")
|
||||
// }
|
||||
|
||||
// if pageAndSort.Page < 0 || pageAndSort.PageSize == 0 {
|
||||
// return errors.New("page and PageSize must not nil")
|
||||
// }
|
||||
|
||||
// skip := pageAndSort.Page * pageAndSort.PageSize
|
||||
// opt.Skip = &skip
|
||||
// opt.Limit = &pageAndSort.PageSize
|
||||
// if pageAndSort.Sort != nil {
|
||||
// opt.Sort = pageAndSort.Sort
|
||||
// }
|
||||
// opts = append(opts, opt)
|
||||
// cur, err := b.collection().Find(ctx, query, opts...)
|
||||
// if err != nil && err == mongo.ErrNoDocuments {
|
||||
// return nil
|
||||
// } else if err != nil {
|
||||
// return err
|
||||
// }
|
||||
// err = cur.All(ctx, bind)
|
||||
// return err
|
||||
|
||||
// }
|
||||
|
||||
func (b *Collection) FindOneAndDelete(ctx context.Context, bind interface{}, query bson.M, opts ...*options.FindOneAndDeleteOptions) error {
|
||||
result := b.coll.FindOneAndDelete(ctx, query, opts...)
|
||||
if result.Err() != nil {
|
||||
return result.Err()
|
||||
}
|
||||
return result.Decode(bind)
|
||||
}
|
||||
|
||||
func (b *Collection) InsertOneTiny(ctx context.Context, doc interface{}) (err error) {
|
||||
_, err = b.coll.InsertOne(ctx, doc)
|
||||
return
|
||||
}
|
||||
|
||||
func (b *Collection) InsertOne(ctx context.Context, doc interface{}) (int, error) {
|
||||
result, err := b.coll.InsertOne(ctx, doc)
|
||||
if err != nil {
|
||||
// log.Info("Collection:InsertOne failed", log.Any("colName", b.name), log.E(err))
|
||||
return 0, err
|
||||
}
|
||||
if _, ok := result.InsertedID.(primitive.ObjectID); ok {
|
||||
return 1, err
|
||||
}
|
||||
return 0, err
|
||||
}
|
||||
func (b *Collection) IDInsertOne(ctx context.Context, doc interface{}) (string, error) {
|
||||
result, err := b.coll.InsertOne(ctx, doc)
|
||||
if err != nil {
|
||||
log.Info("Collection:InsertOne failed", log.Any("colName", b.name), log.E(err))
|
||||
return "", err
|
||||
}
|
||||
if id, ok := result.InsertedID.(primitive.ObjectID); ok {
|
||||
return id.Hex(), err
|
||||
}
|
||||
return "", err
|
||||
}
|
||||
|
||||
func (b *Collection) validSlice(bind interface{}) error {
|
||||
if reflect.TypeOf(bind).Kind() != reflect.Ptr {
|
||||
return errors.New("interface must be a Ptr")
|
||||
}
|
||||
if reflect.TypeOf(bind).Elem().Kind() != reflect.Slice {
|
||||
return errors.New("interface must be a slice")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (b *Collection) InsertMany(ctx context.Context, docs interface{}) ([]interface{}, error) {
|
||||
err := b.validSlice(docs)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
res := b.setTimeForSlice(docs)
|
||||
result, err := b.coll.InsertMany(ctx, res)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return result.InsertedIDs, err
|
||||
}
|
||||
|
||||
// struct方式
|
||||
func (b *Collection) UpdateMOne(ctx context.Context, query bson.M, update interface{}) (*mongo.UpdateResult, error) {
|
||||
um := utils.Struct2Map(update)
|
||||
um["update_at"] = time.Now()
|
||||
set := bson.M{
|
||||
"$set": um,
|
||||
}
|
||||
return b.coll.UpdateOne(ctx, query, set)
|
||||
}
|
||||
|
||||
func (b *Collection) UpdateMOneInc(ctx context.Context, query bson.M, update interface{}, inc interface{}) (*mongo.UpdateResult, error) {
|
||||
um := utils.Struct2Map(update)
|
||||
um["update_at"] = time.Now()
|
||||
set := bson.M{
|
||||
"$set": um,
|
||||
}
|
||||
im := utils.Struct2Map(inc)
|
||||
if im != nil {
|
||||
set["$inc"] = im
|
||||
}
|
||||
return b.coll.UpdateOne(ctx, query, set)
|
||||
}
|
||||
|
||||
// 原生bson.M 方式
|
||||
func (b *Collection) UpdateOne(ctx context.Context, query bson.M, update bson.M, opts ...*options.UpdateOptions) (*mongo.UpdateResult, error) {
|
||||
//b.setUpdatedAt(update, time.Now())
|
||||
return b.coll.UpdateOne(ctx, query, update, opts...)
|
||||
}
|
||||
|
||||
func (b *Collection) UpsertOne(ctx context.Context, query bson.M, update bson.M) (*mongo.UpdateResult, error) {
|
||||
opt := &options.UpdateOptions{}
|
||||
opt.SetUpsert(true)
|
||||
//b.setDeaultTime(update)
|
||||
result, err := b.coll.UpdateOne(ctx, query, update, opt)
|
||||
return result, err
|
||||
|
||||
}
|
||||
|
||||
func (b *Collection) UpdateMMany(ctx context.Context, query bson.M, update interface{}) (*mongo.UpdateResult, error) {
|
||||
// b.setUpdatedAt(update, time.Now())
|
||||
return b.coll.UpdateMany(ctx, query, update)
|
||||
}
|
||||
|
||||
func (b *Collection) UpdateMany(ctx context.Context, query bson.M, update bson.M) (*mongo.UpdateResult, error) {
|
||||
//b.setUpdatedAt(update, time.Now())
|
||||
return b.coll.UpdateMany(ctx, query, update)
|
||||
}
|
||||
|
||||
// func (b *Collection) DeleteById(ctx context.Context, id interface{}) (int64, error) {
|
||||
|
||||
// param := bson.M{"_id": id}
|
||||
// result, err := b.collection().DeleteOne(ctx, param)
|
||||
// if err != nil {
|
||||
// return 0, err
|
||||
// }
|
||||
// return result.DeletedCount, nil
|
||||
|
||||
// }
|
||||
|
||||
func (b *Collection) DeleteOne(ctx context.Context, query bson.M, opts ...*options.DeleteOptions) (int64, error) {
|
||||
result, err := b.coll.DeleteOne(ctx, query)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
return result.DeletedCount, nil
|
||||
}
|
||||
|
||||
func (b *Collection) DeleteMany(ctx context.Context, query bson.M, opts ...*options.DeleteOptions) (int64, error) {
|
||||
result, err := b.coll.DeleteMany(ctx, query, opts...)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
deletedCount := result.DeletedCount
|
||||
return deletedCount, nil
|
||||
|
||||
}
|
||||
|
||||
func (b *Collection) Distince(ctx context.Context, filter bson.M, fieldName string) ([]interface{}, error) {
|
||||
return b.coll.Distinct(ctx, fieldName, filter)
|
||||
}
|
||||
|
||||
func (b *Collection) Aggregate(ctx context.Context, bind interface{}, pipeline []bson.M, opts ...*options.AggregateOptions) error {
|
||||
|
||||
err := b.validSlice(bind)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
//bs, _ := json.Marshal(pipeline)
|
||||
//log.Debug(string(bs))
|
||||
cur, err := b.coll.Aggregate(ctx, pipeline, opts...)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return cur.All(ctx, bind)
|
||||
}
|
||||
|
||||
func (b *Collection) ObjectAggregate(ctx context.Context, bind interface{}, pipeline []bson.M, opts ...*options.AggregateOptions) error {
|
||||
cur, err := b.coll.Aggregate(ctx, pipeline, opts...)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if cur.Next(ctx) {
|
||||
return cur.Decode(bind)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (b *Collection) QueryOne(ctx context.Context, bind interface{}, query interface{}, opts ...*options.FindOneOptions) error {
|
||||
err := b.coll.FindOne(ctx, query, opts...).Decode(bind)
|
||||
if err != nil && err == mongo.ErrNoDocuments {
|
||||
return nil
|
||||
}
|
||||
return err
|
||||
}
|
||||
|
||||
func (b *Collection) Count(ctx context.Context, query interface{}, opts ...*options.CountOptions) (int64, error) {
|
||||
return b.coll.CountDocuments(ctx, query, opts...)
|
||||
}
|
||||
|
||||
func (b *Collection) EstimatedCount(ctx context.Context, opts ...*options.EstimatedDocumentCountOptions) (int64, error) {
|
||||
return b.coll.EstimatedDocumentCount(ctx, opts...)
|
||||
}
|
||||
|
||||
func (b *Collection) BulkWrite(ctx context.Context, models []mongo.WriteModel, opts ...*options.BulkWriteOptions) (*mongo.BulkWriteResult, error) {
|
||||
return b.coll.BulkWrite(ctx, models, opts...)
|
||||
}
|
||||
|
||||
// CreateIndex 创建数据索引.options 在index里面创建
|
||||
func (b *Collection) CreateIndex(ctx context.Context, models []mongo.IndexModel) ([]string, error) {
|
||||
data, err := b.coll.Indexes().CreateMany(ctx, models)
|
||||
if err != nil {
|
||||
if strings.Contains(err.Error(), "IndexKeySpecsConflict") {
|
||||
return data, nil
|
||||
}
|
||||
if strings.Contains(err.Error(), "IndexOptionsConflict") {
|
||||
return data, nil
|
||||
}
|
||||
return data, err
|
||||
}
|
||||
return data, nil
|
||||
}
|
||||
|
||||
// DropIndex 删除数据索引.options 在index里面创建
|
||||
// func (b *Collection) DropIndex(ctx context.Context, indexName string) error {
|
||||
// _, err := b.coll.Indexes().DropOne(ctx, indexName)
|
||||
// if err != nil {
|
||||
// return err
|
||||
// }
|
||||
// return nil
|
||||
// }
|
||||
|
||||
// func (b *MongoTool) addCreatedAt(val interface{}) {
|
||||
// if _, ok := val.(bson.M); ok {
|
||||
// return
|
||||
// }
|
||||
// if _, ok := val.(map[string]interface{}); ok {
|
||||
// return
|
||||
// }
|
||||
// v := reflect.ValueOf(val).Elem().FieldByName("CreatedAt")
|
||||
// if v.IsValid() {
|
||||
// v.Set(reflect.ValueOf(time.Now()))
|
||||
// }
|
||||
|
||||
// }
|
||||
// func (b *MongoTool) addUpdatedAt(val interface{}) {
|
||||
// if _, ok := val.(bson.M); ok {
|
||||
// return
|
||||
// }
|
||||
// if _, ok := val.(map[string]interface{}); ok {
|
||||
// return
|
||||
// }
|
||||
// v := reflect.ValueOf(val).Elem().FieldByName("UpdatedAt")
|
||||
// if v.IsValid() {
|
||||
// v.Set(reflect.ValueOf(time.Now()))
|
||||
// }
|
||||
// }
|
||||
|
||||
func (b *Collection) setTimeForSlice(docs interface{}) []interface{} {
|
||||
arr := reflect.ValueOf(docs).Elem()
|
||||
result := make([]interface{}, 0)
|
||||
for i := 0; i < arr.Len(); i++ {
|
||||
ele := arr.Index(i).Elem()
|
||||
now := time.Now()
|
||||
if ma := ele.FieldByName("update_at"); ma.IsValid() {
|
||||
ma.Set(reflect.ValueOf(now))
|
||||
}
|
||||
if ca := ele.FieldByName("create_at"); ca.IsValid() {
|
||||
ca.Set(reflect.ValueOf(now))
|
||||
}
|
||||
result = append(result, ele.Interface())
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
// func (b *Collection) validSlice(bind interface{}) error {
|
||||
|
||||
// if reflect.TypeOf(bind).Kind() != reflect.Ptr {
|
||||
// return errors.New("interface must be a Ptr")
|
||||
// }
|
||||
// if reflect.TypeOf(bind).Elem().Kind() != reflect.Slice {
|
||||
// return errors.New("interface must be a slice")
|
||||
// }
|
||||
// return nil
|
||||
// }
|
||||
|
||||
// func (b *Collection) setUpdatedAt(doc bson.M, now time.Time) {
|
||||
// set, exists := doc["$set"]
|
||||
// if !exists {
|
||||
// doc["$set"] = bson.M{"updatedAt": now}
|
||||
// return
|
||||
// }
|
||||
|
||||
// if res, ok := set.(bson.M); ok {
|
||||
// res["updatedAt"] = now
|
||||
// }
|
||||
// }
|
||||
|
||||
// func (b *Collection) setCreatedAt(doc bson.M, now time.Time) {
|
||||
// set, exists := doc["$set"]
|
||||
// if !exists {
|
||||
// doc["$set"] = bson.M{"createdAt": now}
|
||||
// return
|
||||
// }
|
||||
// if res, ok := set.(bson.M); ok {
|
||||
// res["createdAt"] = now
|
||||
// }
|
||||
// }
|
||||
|
||||
// func (b *Collection) setDeaultTime(doc bson.M) {
|
||||
// now := time.Now()
|
||||
// set, exists := doc["$set"]
|
||||
// if exists {
|
||||
// if res, ok := set.(bson.M); ok {
|
||||
// delete(res, "createdAt")
|
||||
// res["updatedAt"] = now
|
||||
// if d, ok := doc["$setOnInsert"]; ok {
|
||||
// if ss, ok := d.(bson.M); ok {
|
||||
// if _, ok := ss["createdAt"]; !ok {
|
||||
// ss["createdAt"] = now
|
||||
// }
|
||||
// }
|
||||
// } else {
|
||||
// doc["$setOnInsert"] = bson.M{"createdAt": now}
|
||||
// }
|
||||
// return
|
||||
// }
|
||||
// if res, ok := set.(interface{}); ok {
|
||||
// struct2Map := utils.Struct2Map(res)
|
||||
// delete(struct2Map, "createdAt")
|
||||
// struct2Map["updatedAt"] = now
|
||||
// doc["$set"] = struct2Map
|
||||
// doc["$setOnInsert"] = bson.M{"createdAt": now}
|
||||
// return
|
||||
// }
|
||||
// }
|
||||
// }
|
||||
|
||||
// ////////////////////////////////////////////
|
||||
func (b *Collection) GetStringIDInsertOne(ctx context.Context, doc interface{}) (string, error) {
|
||||
result, err := b.coll.InsertOne(ctx, doc)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if id, ok := result.InsertedID.(primitive.ObjectID); ok {
|
||||
return id.Hex(), nil
|
||||
}
|
||||
return "", err
|
||||
}
|
||||
|
||||
// InsertOne 插入单条信息
|
||||
func (b *Collection) GetObjectIDInsertOne(ctx context.Context, document interface{}, opts ...*options.InsertOneOptions) (*mongo.InsertOneResult, error) {
|
||||
return b.coll.InsertOne(ctx, document, opts...)
|
||||
}
|
||||
|
||||
// UpdateOneNew 或者修改或者插入一条数据
|
||||
func (b *Collection) UpdateOneNew(ctx context.Context, query bson.M, update bson.M) (*mongo.UpdateResult, error) {
|
||||
return b.coll.UpdateOne(ctx, query, update, options.Update().SetUpsert(true))
|
||||
}
|
||||
|
||||
func (b *Collection) Max(ctx context.Context, fieldName string, filter bson.M) (int64, error) {
|
||||
//b.coll.Aggregate(ctx, )
|
||||
type Temp struct {
|
||||
Max int64 `bson:"max"`
|
||||
}
|
||||
var temp []Temp
|
||||
query := []bson.M{
|
||||
{"$group": bson.M{
|
||||
"_id": nil,
|
||||
"max": bson.M{"$max": "$" + fieldName},
|
||||
},
|
||||
},
|
||||
}
|
||||
if err := b.Aggregate(ctx, &temp, query); err != nil || len(temp) <= 0 {
|
||||
return 0, err
|
||||
}
|
||||
|
||||
return temp[0].Max, nil
|
||||
}
|
||||
|
||||
// func (b *Collection) setUpdatedAt(doc bson.M, now time.Time) {
|
||||
// set, exists := doc["$set"]
|
||||
// if !exists {
|
||||
// doc["$set"] = bson.M{"updatedAt": now}
|
||||
// return
|
||||
// }
|
||||
|
||||
// if res, ok := set.(bson.M); ok {
|
||||
// res["updatedAt"] = now
|
||||
// }
|
||||
// }
|
||||
|
||||
func (b *Collection) FindOneAndUpdate(ctx context.Context, bind interface{}, query bson.M, update bson.M, opts ...*options.FindOneAndUpdateOptions) error {
|
||||
setReturn := false
|
||||
for _, opt := range opts {
|
||||
if opt != nil && opt.ReturnDocument != nil {
|
||||
setReturn = true
|
||||
break
|
||||
}
|
||||
}
|
||||
if !setReturn {
|
||||
after := options.After
|
||||
opts = append(opts, &options.FindOneAndUpdateOptions{ReturnDocument: &after})
|
||||
}
|
||||
// b.coll.setUpdatedAt(update, time.Now())
|
||||
result := b.coll.FindOneAndUpdate(ctx, query, update, opts...)
|
||||
err := result.Decode(bind)
|
||||
return err
|
||||
}
|
||||
|
||||
// //
|
||||
// func (b *Collection) IfNotExistThenInsertAndUpdate(ctx context.Context, query bson.M, insert bson.M, update bson.M) error {
|
||||
// // func (b *MongoTool) FindOneAndUpdate(ctx context.Context, bind interface{}, query bson.M, update bson.M, opts ...*options.FindOneAndUpdateOptions) error {
|
||||
// err := b.FindOneAndUpdate(ctx, nil, query, update)
|
||||
// if err != nil {
|
||||
// _, err = b.InsertOne(ctx, insert)
|
||||
// }
|
||||
// return err
|
||||
// }
|
||||
|
||||
func (b *Collection) WatchOpLog(ctx context.Context, pipeline []bson.M,
|
||||
opt ...*options.ChangeStreamOptions) (*mongo.ChangeStream, error) {
|
||||
changeStream, err := b.coll.Watch(ctx, pipeline, opt...)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return changeStream, nil
|
||||
}
|
||||
@@ -0,0 +1,119 @@
|
||||
package mongos
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"github.com/shopspring/decimal"
|
||||
"go.mongodb.org/mongo-driver/bson"
|
||||
"go.mongodb.org/mongo-driver/bson/bsoncodec"
|
||||
"go.mongodb.org/mongo-driver/bson/bsonrw"
|
||||
"go.mongodb.org/mongo-driver/bson/bsontype"
|
||||
"go.mongodb.org/mongo-driver/bson/primitive"
|
||||
"reflect"
|
||||
"time"
|
||||
)
|
||||
|
||||
var Registry = func() *bsoncodec.Registry {
|
||||
builder := bson.NewRegistryBuilder()
|
||||
builder.RegisterDecoder(reflect.TypeOf(time.Time{}), &localTimeDecoder{})
|
||||
builder.RegisterDecoder(reflect.TypeOf(decimal.Decimal{}), &decimalDecoder{})
|
||||
builder.RegisterEncoder(reflect.TypeOf(decimal.Decimal{}), &decimalEncoder{})
|
||||
builder.RegisterDefaultDecoder(reflect.Float32, &floatDecoder{})
|
||||
builder.RegisterDefaultDecoder(reflect.Float64, &floatDecoder{})
|
||||
return builder.Build()
|
||||
}()
|
||||
|
||||
type floatDecoder struct {
|
||||
}
|
||||
|
||||
func (dvd *floatDecoder) DecodeValue(dc bsoncodec.DecodeContext, vr bsonrw.ValueReader, val reflect.Value) error {
|
||||
var f float64
|
||||
var err error
|
||||
switch vr.Type() {
|
||||
case bsontype.Int32:
|
||||
i32, err := vr.ReadInt32()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
f = float64(i32)
|
||||
case bsontype.Int64:
|
||||
i64, err := vr.ReadInt64()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
f = float64(i64)
|
||||
case bsontype.Double:
|
||||
f, err = vr.ReadDouble()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
default:
|
||||
return fmt.Errorf("cannot decode %v into a float32 or float64 type", vr.Type())
|
||||
}
|
||||
val.SetFloat(f)
|
||||
return nil
|
||||
}
|
||||
|
||||
type localTimeDecoder struct{}
|
||||
|
||||
func (*localTimeDecoder) DecodeValue(dc bsoncodec.DecodeContext, vr bsonrw.ValueReader, val reflect.Value) error {
|
||||
dvd := bsoncodec.DefaultValueDecoders{}
|
||||
err := dvd.TimeDecodeValue(dc, vr, val)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
t := val.Interface().(time.Time)
|
||||
val.Set(reflect.ValueOf(t.Local()))
|
||||
return nil
|
||||
}
|
||||
|
||||
type decimalDecoder struct{}
|
||||
|
||||
func (*decimalDecoder) DecodeValue(dc bsoncodec.DecodeContext, vr bsonrw.ValueReader, val reflect.Value) error {
|
||||
if !val.IsValid() || val.Type() != reflect.TypeOf(decimal.Decimal{}) {
|
||||
return bsoncodec.ValueDecoderError{Name: "DecimalDecodeValue", Types: []reflect.Type{reflect.TypeOf(decimal.Decimal{})}, Received: val}
|
||||
}
|
||||
var (
|
||||
num bool
|
||||
)
|
||||
if vr.Type() == bson.TypeInt32 || vr.Type() == bson.TypeInt64 {
|
||||
vr.ReadInt32()
|
||||
vr.ReadInt64()
|
||||
num = true
|
||||
}
|
||||
if num {
|
||||
d := decimal.NewFromFloat(0.0)
|
||||
val.Set(reflect.ValueOf(d))
|
||||
return nil
|
||||
}
|
||||
|
||||
mongodecimal, err := vr.ReadDecimal128()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
d, err := decimal.NewFromString(mongodecimal.String())
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
val.Set(reflect.ValueOf(d))
|
||||
return nil
|
||||
}
|
||||
|
||||
//
|
||||
|
||||
type decimalEncoder struct{}
|
||||
|
||||
func (*decimalEncoder) EncodeValue(ctx bsoncodec.EncodeContext, vw bsonrw.ValueWriter, val reflect.Value) error {
|
||||
if !val.IsValid() || val.Type() != reflect.TypeOf(decimal.Decimal{}) {
|
||||
return bsoncodec.ValueDecoderError{Name: "DecimalEncodeValue", Types: []reflect.Type{reflect.TypeOf(decimal.Decimal{})}, Received: val}
|
||||
}
|
||||
|
||||
if d, ok := val.Interface().(decimal.Decimal); ok {
|
||||
mongodecimal, err := primitive.ParseDecimal128(d.StringFixed(2))
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
val = reflect.ValueOf(mongodecimal)
|
||||
}
|
||||
dve := bsoncodec.DefaultValueEncoders{}
|
||||
return dve.Decimal128EncodeValue(ctx, vw, val)
|
||||
}
|
||||
@@ -0,0 +1,142 @@
|
||||
package mongos
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"videoweb/pub/models/utils"
|
||||
|
||||
"go.mongodb.org/mongo-driver/bson"
|
||||
"go.mongodb.org/mongo-driver/mongo"
|
||||
"go.mongodb.org/mongo-driver/mongo/options"
|
||||
)
|
||||
|
||||
var Lid struct {
|
||||
User *lid
|
||||
MwAcc *lid
|
||||
MwRole *lid
|
||||
MwOpLog *lid
|
||||
MwAccToken *lid
|
||||
SystemMsg *lid
|
||||
AdvitiseInfo *lid
|
||||
AdvitiseBind *lid
|
||||
AdvitiseLocation *lid
|
||||
App *lid
|
||||
VideoTag *lid
|
||||
VideoCategory *lid
|
||||
VideoTopic *lid
|
||||
VideoInfo *lid
|
||||
VideoLike *lid
|
||||
VideoCollect *lid
|
||||
VideoWatchHistory *lid
|
||||
TransCategory *lid
|
||||
Translation *lid
|
||||
VideoCreator *lid
|
||||
Comment *lid
|
||||
HtmlText *lid
|
||||
Actress *lid
|
||||
Post *lid
|
||||
PostCategory *lid
|
||||
Card *lid
|
||||
CoinRecharge *lid
|
||||
RechargePackage *lid
|
||||
ContactUs *lid
|
||||
Feedback *lid
|
||||
VipOrder *lid
|
||||
CoinTrans *lid
|
||||
Playlet *lid
|
||||
PlayletChapter *lid
|
||||
PlayletChapterLike *lid
|
||||
PlayletChapterCollect *lid
|
||||
PlayletChapterUnlock *lid
|
||||
PlayletTheme *lid
|
||||
PlayletTopicGroup *lid
|
||||
PlayletTag *lid
|
||||
PlayletLike *lid
|
||||
PlayletCollect *lid
|
||||
PlayletWatchProgress *lid
|
||||
PlayletWatchHistory *lid
|
||||
CommentLike *lid
|
||||
Announcement *lid
|
||||
HotSearch *lid
|
||||
HomeBanner *lid
|
||||
HomeConfig *lid
|
||||
}
|
||||
|
||||
func StartLid() {
|
||||
|
||||
coll := Shop.Coll(utils.TN.GenerateId)
|
||||
initIndex(coll)
|
||||
|
||||
Lid.User = createLid(coll, "user", 10)
|
||||
Lid.MwAcc = createLid(coll, "MwAcc", 1)
|
||||
Lid.MwRole = createLid(coll, "MwRole", 1)
|
||||
Lid.MwOpLog = createLid(coll, "MwOpLog", 10)
|
||||
Lid.MwAccToken = createLid(coll, "MwAccToken", 10)
|
||||
Lid.SystemMsg = createLid(coll, "SystemMsg", 10)
|
||||
Lid.AdvitiseInfo = createLid(coll, "AdvitiseInfo", 10)
|
||||
Lid.AdvitiseBind = createLid(coll, "AdvitiseBind", 10)
|
||||
Lid.AdvitiseLocation = createLid(coll, "AdvitiseLocation", 10)
|
||||
Lid.App = createLid(coll, "App", 10)
|
||||
Lid.VideoTag = createLid(coll, "VideoTag", 10)
|
||||
Lid.VideoCategory = createLid(coll, "VideoCategory", 10)
|
||||
Lid.VideoInfo = createLid(coll, "VideoInfo", 10)
|
||||
Lid.VideoLike = createLid(coll, "videoLike", 10)
|
||||
Lid.VideoCollect = createLid(coll, "videoCollect", 10)
|
||||
Lid.VideoWatchHistory = createLid(coll, "videoWatchHistory", 10)
|
||||
Lid.TransCategory = createLid(coll, "TransCategory", 10)
|
||||
Lid.Translation = createLid(coll, "TranslationData", 10)
|
||||
Lid.VideoCreator = createLid(coll, "VideoCreator", 10)
|
||||
Lid.Comment = createLid(coll, "Comment", 10)
|
||||
Lid.HtmlText = createLid(coll, "HtmlText", 10)
|
||||
Lid.Actress = createLid(coll, "Actress", 10)
|
||||
Lid.Post = createLid(coll, "Post", 10)
|
||||
Lid.VideoTopic = createLid(coll, "VideoTopic", 10)
|
||||
Lid.PostCategory = createLid(coll, "PostCategory", 10)
|
||||
Lid.Card = createLid(coll, "card", 10)
|
||||
Lid.CoinRecharge = createLid(coll, "coinRecharge", 10)
|
||||
Lid.RechargePackage = createLid(coll, "rechargePackage", 10)
|
||||
Lid.ContactUs = createLid(coll, "contactus", 10)
|
||||
Lid.VipOrder = createLid(coll, "vipOrder", 10)
|
||||
Lid.CoinTrans = createLid(coll, "coinTrans", 10)
|
||||
Lid.Feedback = createLid(coll, "feedback", 10)
|
||||
Lid.Playlet = createLid(coll, "playlet", 10)
|
||||
Lid.PlayletChapter = createLid(coll, "playletChapter", 10)
|
||||
Lid.PlayletChapterLike = createLid(coll, "playletChapterLike", 10)
|
||||
Lid.PlayletChapterCollect = createLid(coll, "playletChapterCollect", 10)
|
||||
Lid.PlayletChapterUnlock = createLid(coll, "playletChapterUnlock", 10)
|
||||
Lid.PlayletTheme = createLid(coll, "playletTheme", 10)
|
||||
Lid.PlayletTopicGroup = createLid(coll, "playletTopicGroup", 10)
|
||||
Lid.PlayletTag = createLid(coll, "playletTag", 10)
|
||||
Lid.PlayletLike = createLid(coll, "playletLike", 10)
|
||||
Lid.PlayletCollect = createLid(coll, "playletCollect", 10)
|
||||
Lid.PlayletWatchProgress = createLid(coll, "playletWatchProgress", 10)
|
||||
Lid.PlayletWatchHistory = createLid(coll, "playletWatchHistory", 10)
|
||||
Lid.CommentLike = createLid(coll, "commentLike", 10)
|
||||
Lid.Announcement = createLid(coll, "announcement", 10)
|
||||
Lid.HotSearch = createLid(coll, "hotSearch", 10)
|
||||
Lid.HomeBanner = createLid(coll, "homeBanner", 10)
|
||||
Lid.HomeConfig = createLid(coll, "homeConfig", 10)
|
||||
}
|
||||
|
||||
// initIndex 初始化索引
|
||||
func initIndex(coll *Collection) {
|
||||
indexes := []mongo.IndexModel{
|
||||
{
|
||||
Keys: bson.D{
|
||||
{Key: utils.FN.GenerateId.BusinessId, Value: 1},
|
||||
},
|
||||
Options: options.Index().SetUnique(true),
|
||||
},
|
||||
}
|
||||
if _, err := coll.CreateIndex(context.Background(), indexes); err != nil {
|
||||
panic(fmt.Sprintf("(user:initIndex) Create Index user failed! err:%+v:", err))
|
||||
}
|
||||
}
|
||||
|
||||
func createLid(db *Collection, businessId string, len int) *lid {
|
||||
lid, err := newLid(db, businessId, len)
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
return lid
|
||||
}
|
||||
@@ -0,0 +1,102 @@
|
||||
package mongos
|
||||
|
||||
import (
|
||||
"context"
|
||||
"time"
|
||||
"videoweb/pub/models/utils"
|
||||
|
||||
"go.mongodb.org/mongo-driver/bson"
|
||||
"go.mongodb.org/mongo-driver/mongo"
|
||||
)
|
||||
|
||||
const notFindMsg = "mongo: no documents in result"
|
||||
|
||||
type lid struct {
|
||||
db *Collection // 数据库连接
|
||||
businessId string // 业务id
|
||||
ch chan int64 // id缓冲池
|
||||
min, max int64 // id段最小值，最大值
|
||||
}
|
||||
|
||||
// newLid 创建一个lid,db:数据库连接；businessId：业务id;len：缓冲池大小
|
||||
func newLid(db *Collection, businessId string, len int) (*lid, error) {
|
||||
lid := lid{
|
||||
db: db,
|
||||
businessId: businessId,
|
||||
ch: make(chan int64, len),
|
||||
}
|
||||
go lid.productId()
|
||||
return &lid, nil
|
||||
}
|
||||
|
||||
// Get 获取自增id
|
||||
func (l *lid) Get() int64 {
|
||||
return <-l.ch
|
||||
}
|
||||
|
||||
// productId 生产id
|
||||
func (l *lid) productId() {
|
||||
l.reset()
|
||||
for {
|
||||
if l.min >= l.max {
|
||||
l.reset()
|
||||
|
||||
}
|
||||
|
||||
l.min++
|
||||
l.ch <- l.min
|
||||
}
|
||||
}
|
||||
|
||||
// reset 不断尝试从数据库获取，直到成功
|
||||
func (l *lid) reset() {
|
||||
for {
|
||||
err := l.getFromDB()
|
||||
if err == nil {
|
||||
return
|
||||
}
|
||||
time.Sleep(time.Second)
|
||||
continue
|
||||
}
|
||||
}
|
||||
|
||||
func (l *lid) getFromDB() error {
|
||||
var (
|
||||
max_id int64
|
||||
step int64
|
||||
//idModel IdModel
|
||||
idModel utils.GenerateId
|
||||
)
|
||||
businessIdfilter := bson.M{utils.FN.GenerateId.BusinessId: l.businessId}
|
||||
txErr := Shop.Transaction(context.Background(), func(sctx mongo.SessionContext) (err error) {
|
||||
err = l.db.FindOneEx(sctx, &idModel, businessIdfilter)
|
||||
if err != nil {
|
||||
if err == mongo.ErrNoDocuments {
|
||||
idModel.MaxId = 10000
|
||||
_, err = l.db.InsertOne(sctx,
|
||||
bson.M{
|
||||
utils.FN.GenerateId.BusinessId: l.businessId,
|
||||
utils.FN.GenerateId.MaxId: 0,
|
||||
utils.FN.GenerateId.Step: 10,
|
||||
utils.FN.GenerateId.Desc: "id自增器",
|
||||
utils.FN.GenerateId.CreateAt: time.Now(),
|
||||
utils.FN.GenerateId.UpdateAt: time.Now()})
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
} else {
|
||||
return
|
||||
}
|
||||
}
|
||||
max_id = idModel.MaxId
|
||||
step = idModel.Step
|
||||
_, err = l.db.UpdateOne(sctx, businessIdfilter, bson.M{"$set": bson.M{utils.FN.GenerateId.MaxId: max_id + step}})
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
l.min = max_id
|
||||
l.max = max_id + step
|
||||
return
|
||||
})
|
||||
return txErr
|
||||
}
|
||||
@@ -0,0 +1,168 @@
|
||||
package mongos
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"time"
|
||||
"videoweb/pub/common/log"
|
||||
|
||||
"go.mongodb.org/mongo-driver/mongo/readconcern"
|
||||
"go.mongodb.org/mongo-driver/mongo/readpref"
|
||||
"go.mongodb.org/mongo-driver/mongo/writeconcern"
|
||||
"go.mongodb.org/mongo-driver/x/mongo/driver/connstring"
|
||||
|
||||
"go.mongodb.org/mongo-driver/mongo"
|
||||
"go.mongodb.org/mongo-driver/mongo/options"
|
||||
)
|
||||
|
||||
// var (
|
||||
// connPools *mongo.Client
|
||||
// dataBase *mongo.Database
|
||||
// )
|
||||
|
||||
var (
|
||||
//业务库
|
||||
Shop *Manager
|
||||
//日志库
|
||||
ShopLog *Manager
|
||||
)
|
||||
|
||||
type Manager struct {
|
||||
client *mongo.Client
|
||||
//dataBase *mongo.Database
|
||||
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) Coll(name string) *Collection {
|
||||
coll := m.Database().Collection(name)
|
||||
return &Collection{name: name, coll: coll}
|
||||
}
|
||||
|
||||
func (m *Manager) Init() error {
|
||||
cs, err := connstring.Parse(m.url)
|
||||
if err != nil {
|
||||
log.Error("db connect parse err", log.E(err))
|
||||
return err
|
||||
}
|
||||
if cs.Database == "" {
|
||||
log.Error("db connect database MUST contains")
|
||||
return errors.New("database MUST contains")
|
||||
}
|
||||
|
||||
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
|
||||
}
|
||||
|
||||
//connPools = client
|
||||
m.client = client
|
||||
m.db = cs.Database
|
||||
|
||||
log.Info("mongos:Init ok!", log.Any("dbName", m.db))
|
||||
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 切记里面的context必须要使用同一个，因为里面有重试机制，循环调用该函数
|
||||
func (m *Manager) Transaction(ctx context.Context, fn func(sctx mongo.SessionContext) error) error {
|
||||
opts := options.Session().SetDefaultReadPreference(readpref.Primary())
|
||||
handler := func(sctx mongo.SessionContext) error {
|
||||
lHandler := func(ctx mongo.SessionContext) (err error) {
|
||||
tOpts := options.Transaction()
|
||||
tOpts.SetReadConcern(readconcern.Snapshot())
|
||||
tOpts.SetWriteConcern(writeconcern.New(writeconcern.WMajority()))
|
||||
err = sctx.StartTransaction(tOpts)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
defer func() {
|
||||
if rErr := recover(); rErr != nil {
|
||||
log.Info("caught panic during transaction, aborting...", log.Any("", rErr))
|
||||
aErr := sctx.AbortTransaction(sctx)
|
||||
if aErr != nil {
|
||||
log.Info("中止事务错误", log.E(aErr))
|
||||
}
|
||||
}
|
||||
}()
|
||||
err = fn(sctx)
|
||||
if err != nil {
|
||||
aErr := sctx.AbortTransaction(sctx)
|
||||
if aErr != nil {
|
||||
log.Info("中止事务错误", log.E(aErr))
|
||||
}
|
||||
log.Info("事务内错误", log.E(err))
|
||||
return
|
||||
}
|
||||
return commitWithRetry(sctx)
|
||||
}
|
||||
return runTransactionWithRetry(sctx, lHandler)
|
||||
}
|
||||
return m.client.UseSessionWithOptions(ctx, opts, handler)
|
||||
}
|
||||
|
||||
func runTransactionWithRetry(sctx mongo.SessionContext, txnFn func(mongo.SessionContext) error) error {
|
||||
var err error
|
||||
for i := int64(0); i < 100; i++ {
|
||||
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
|
||||
}
|
||||
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 connMgr.Transaction(ctx, fn)
|
||||
// }
|
||||
Reference in New Issue
Block a user