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 }