Initial commit
This commit is contained in:
@@ -0,0 +1,507 @@
|
||||
package mongos
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"reflect"
|
||||
"stat/common/utils"
|
||||
"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 MongoTool struct {
|
||||
db *Manager
|
||||
collName string
|
||||
}
|
||||
|
||||
type PageAndSort struct {
|
||||
Page int64
|
||||
PageSize int64
|
||||
Sort bson.M //bson.M{_id : -1 }
|
||||
}
|
||||
|
||||
func New(db *Manager, collName string) *MongoTool {
|
||||
return &MongoTool{db: db, collName: collName}
|
||||
}
|
||||
|
||||
func (b *MongoTool) GetCollection() *mongo.Collection {
|
||||
return b.collection()
|
||||
}
|
||||
|
||||
func (b *MongoTool) collection() *mongo.Collection {
|
||||
opts := options.CollectionOptions{Registry: Registry}
|
||||
return b.db.Database().Collection(b.collName, &opts)
|
||||
}
|
||||
|
||||
func (b *MongoTool) FindById(ctx context.Context, bind interface{}, id interface{}, opts ...*options.FindOneOptions) error {
|
||||
param := bson.M{"_id": id}
|
||||
return b.FindOne(ctx, bind, param, opts...)
|
||||
}
|
||||
|
||||
func (b *MongoTool) FindOne(ctx context.Context, bind interface{}, query bson.M, opts ...*options.FindOneOptions) error {
|
||||
err := b.collection().FindOne(ctx, query, opts...).Decode(bind)
|
||||
if err != nil && err == mongo.ErrNoDocuments {
|
||||
return nil
|
||||
}
|
||||
return err
|
||||
}
|
||||
|
||||
func (b *MongoTool) FindOneEx(ctx context.Context, bind interface{}, query bson.M, opts ...*options.FindOneOptions) error {
|
||||
err := b.collection().FindOne(ctx, query, opts...).Decode(bind)
|
||||
return err
|
||||
}
|
||||
|
||||
func (b *MongoTool) 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["updatedAt"] = time.Now()
|
||||
set := bson.M{
|
||||
"$set": um,
|
||||
}
|
||||
result := b.collection().FindOneAndUpdate(ctx, query, set, opts...)
|
||||
err := result.Decode(bind)
|
||||
return err
|
||||
}
|
||||
|
||||
func (b *MongoTool) FindMOneAndUpdateInc(ctx context.Context, bind interface{}, query bson.M, update interface{}, inc interface{}, opts ...*options.FindOneAndUpdateOptions) error {
|
||||
um := utils.Struct2Map(update)
|
||||
um["updatedAt"] = 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.collection().FindOneAndUpdate(ctx, query, set, opts...)
|
||||
err := result.Decode(bind)
|
||||
return err
|
||||
}
|
||||
|
||||
func (b *MongoTool) FindOneAndUpdateTiny(
|
||||
ctx context.Context,
|
||||
query bson.M,
|
||||
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 *MongoTool) 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 *MongoTool) Find(ctx context.Context, bind interface{}, query bson.M, opt ...*options.FindOptions) error {
|
||||
err := b.validSlice(bind)
|
||||
cur, err := b.collection().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 *MongoTool) 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 - 1) * 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 *MongoTool) FindOneAndDelete(ctx context.Context, bind interface{}, query bson.M, opts ...*options.FindOneAndDeleteOptions) error {
|
||||
result := b.collection().FindOneAndDelete(ctx, query, opts...)
|
||||
if result.Err() != nil {
|
||||
return result.Err()
|
||||
}
|
||||
return result.Decode(bind)
|
||||
}
|
||||
|
||||
func (b *MongoTool) InsertOneTiny(ctx context.Context, doc interface{}) (err error) {
|
||||
_, err = b.collection().InsertOne(ctx, doc)
|
||||
return
|
||||
}
|
||||
|
||||
func (b *MongoTool) InsertOne(ctx context.Context, doc interface{}) (string, error) {
|
||||
//设置创建时间
|
||||
b.addCreatedAt(doc)
|
||||
//设置修改时间
|
||||
b.addUpdatedAt(doc)
|
||||
result, err := b.collection().InsertOne(ctx, doc)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if id, ok := result.InsertedID.(primitive.ObjectID); ok {
|
||||
return id.Hex(), nil
|
||||
}
|
||||
return "", err
|
||||
}
|
||||
|
||||
func (b *MongoTool) 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.collection().InsertMany(ctx, res)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return result.InsertedIDs, err
|
||||
}
|
||||
|
||||
// struct方式
|
||||
func (b *MongoTool) UpdateMOne(ctx context.Context, query bson.M, update interface{}) (*mongo.UpdateResult, error) {
|
||||
um := utils.Struct2Map(update)
|
||||
um["updatedAt"] = time.Now()
|
||||
set := bson.M{
|
||||
"$set": um,
|
||||
}
|
||||
return b.collection().UpdateOne(ctx, query, set)
|
||||
}
|
||||
|
||||
func (b *MongoTool) UpdateMOneInc(ctx context.Context, query bson.M, update interface{}, inc interface{}) (*mongo.UpdateResult, error) {
|
||||
um := utils.Struct2Map(update)
|
||||
um["updatedAt"] = time.Now()
|
||||
set := bson.M{
|
||||
"$set": um,
|
||||
}
|
||||
im := utils.Struct2Map(inc)
|
||||
if im != nil {
|
||||
set["$inc"] = im
|
||||
}
|
||||
return b.collection().UpdateOne(ctx, query, set)
|
||||
}
|
||||
|
||||
// 原生bson.M 方式
|
||||
func (b *MongoTool) UpdateOne(ctx context.Context, query bson.M, update bson.M) (*mongo.UpdateResult, error) {
|
||||
b.setUpdatedAt(update, time.Now())
|
||||
return b.collection().UpdateOne(ctx, query, update)
|
||||
}
|
||||
|
||||
// struct方式
|
||||
func (b *MongoTool) UpsertMOne(ctx context.Context, query bson.M, update interface{}) (*mongo.UpdateResult, error) {
|
||||
um := utils.Struct2Map(update)
|
||||
um["updatedAt"] = time.Now()
|
||||
set := bson.M{
|
||||
"$set": um,
|
||||
}
|
||||
upsert := true
|
||||
opt := &options.UpdateOptions{
|
||||
Upsert: &upsert,
|
||||
}
|
||||
now := time.Now()
|
||||
result, err := b.collection().UpdateOne(ctx, query, set, opt)
|
||||
if err == nil {
|
||||
if result.UpsertedCount >= 1 {
|
||||
_, err = b.collection().UpdateOne(ctx, query, bson.M{"$set": bson.M{"createdAt": now}})
|
||||
}
|
||||
}
|
||||
return result, err
|
||||
|
||||
}
|
||||
|
||||
func (b *MongoTool) UpsertOneEx(ctx context.Context, query bson.M, update bson.M) (*mongo.UpdateResult, error) {
|
||||
opt := &options.UpdateOptions{}
|
||||
opt.SetUpsert(true)
|
||||
result, err := b.collection().UpdateOne(ctx, query, update, opt)
|
||||
return result, err
|
||||
|
||||
}
|
||||
|
||||
func (b *MongoTool) UpsertOne(ctx context.Context, query bson.M, update bson.M) (*mongo.UpdateResult, error) {
|
||||
|
||||
upsert := true
|
||||
opt := &options.UpdateOptions{
|
||||
Upsert: &upsert,
|
||||
}
|
||||
now := time.Now()
|
||||
b.setUpdatedAt(update, now)
|
||||
result, err := b.collection().UpdateOne(ctx, query, update, opt)
|
||||
if err == nil {
|
||||
if result.UpsertedCount >= 1 {
|
||||
_, err = b.collection().UpdateOne(ctx, query, bson.M{"$set": bson.M{"createdAt": now}})
|
||||
}
|
||||
}
|
||||
return result, err
|
||||
|
||||
}
|
||||
|
||||
func (b *MongoTool) UpdateMMany(ctx context.Context, query bson.M, update interface{}) (*mongo.UpdateResult, error) {
|
||||
//b.setUpdatedAt(update, time.Now())
|
||||
return b.collection().UpdateMany(ctx, query, update)
|
||||
}
|
||||
|
||||
func (b *MongoTool) UpdateMany(ctx context.Context, query bson.M, update bson.M) (*mongo.UpdateResult, error) {
|
||||
b.setUpdatedAt(update, time.Now())
|
||||
return b.collection().UpdateMany(ctx, query, update)
|
||||
}
|
||||
|
||||
func (b *MongoTool) 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 *MongoTool) DeleteOne(ctx context.Context, query bson.M, opts ...*options.DeleteOptions) (int64, error) {
|
||||
|
||||
result, err := b.collection().DeleteOne(ctx, query)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
return result.DeletedCount, nil
|
||||
}
|
||||
|
||||
func (b *MongoTool) DeleteMany(ctx context.Context, query bson.M, opts ...*options.DeleteOptions) (int64, error) {
|
||||
|
||||
result, err := b.collection().DeleteMany(ctx, query, opts...)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
deletedCount := result.DeletedCount
|
||||
return deletedCount, nil
|
||||
}
|
||||
|
||||
func (b *MongoTool) Aggregate(ctx context.Context, bind interface{}, pipeline []bson.M, opts ...*options.AggregateOptions) error {
|
||||
|
||||
err := b.validSlice(bind)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
cur, err := b.collection().Aggregate(ctx, pipeline, opts...)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return cur.All(ctx, bind)
|
||||
}
|
||||
|
||||
func (b *MongoTool) ObjectAggregate(ctx context.Context, bind interface{}, pipeline []bson.M, opts ...*options.AggregateOptions) error {
|
||||
cur, err := b.collection().Aggregate(ctx, pipeline, opts...)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if cur.Next(ctx) {
|
||||
return cur.Decode(bind)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (b *MongoTool) QueryOne(ctx context.Context, bind interface{}, query interface{}, opts ...*options.FindOneOptions) error {
|
||||
err := b.collection().FindOne(ctx, query, opts...).Decode(bind)
|
||||
if err != nil && err == mongo.ErrNoDocuments {
|
||||
return nil
|
||||
}
|
||||
return err
|
||||
}
|
||||
|
||||
func (b *MongoTool) Count(ctx context.Context, query interface{}, opts ...*options.CountOptions) (int64, error) {
|
||||
return b.collection().CountDocuments(ctx, query, opts...)
|
||||
}
|
||||
|
||||
func (b *MongoTool) EstimatedCount(ctx context.Context, opts ...*options.EstimatedDocumentCountOptions) (int64, error) {
|
||||
return b.collection().EstimatedDocumentCount(ctx, opts...)
|
||||
}
|
||||
|
||||
func (b *MongoTool) BulkWrite(ctx context.Context, models []mongo.WriteModel, opts ...*options.BulkWriteOptions) (*mongo.BulkWriteResult, error) {
|
||||
return b.collection().BulkWrite(ctx, models, opts...)
|
||||
}
|
||||
|
||||
// CreateIndex 创建数据索引.options 在index里面创建
|
||||
func (b *MongoTool) CreateIndex(ctx context.Context, models []mongo.IndexModel) ([]string, error) {
|
||||
data, err := b.collection().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 *MongoTool) DropIndex(ctx context.Context, indexName string) error {
|
||||
_, err := b.collection().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 *MongoTool) 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("UpdatedAt"); ma.IsValid() {
|
||||
ma.Set(reflect.ValueOf(now))
|
||||
}
|
||||
if ca := ele.FieldByName("CreatedAt"); ca.IsValid() {
|
||||
ca.Set(reflect.ValueOf(now))
|
||||
}
|
||||
result = append(result, ele.Interface())
|
||||
}
|
||||
return result
|
||||
}
|
||||
func (b *MongoTool) 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 *MongoTool) 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 *MongoTool) 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
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
// Exists 是否存在数据
|
||||
func (b *MongoTool) Exists(ctx context.Context, filter interface{}, opts ...*options.FindOneOptions) (bool, error) {
|
||||
var limit int64 = 1
|
||||
lo := &options.CountOptions{Limit: &limit}
|
||||
n, err := b.Count(ctx, filter, lo)
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
return n > 0, 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,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
|
||||
}
|
||||
Reference in New Issue
Block a user