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 }