package db import ( "context" "errors" "fmt" "reflect" "time" "91porn-server/common/log" "91porn-server/common/stderr" "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" ) var ( enableSortMCheck = true bsonDType = reflect.TypeOf(bson.D{}) bsonDPtrType = reflect.TypeOf(&bson.D{}) bsonEType = reflect.TypeOf(bson.E{}) bsonEPtrType = reflect.TypeOf(&bson.E{}) bsonMType = reflect.TypeOf(bson.M{}) bsonMPtrType = reflect.TypeOf(&bson.M{}) ) var skipErrors = []error{mongo.ErrNoDocuments} func handleDbError(err error) error { for _, e := range skipErrors { if err == e { return nil } } if err != nil { if mongo.IsDuplicateKeyError(err) { return stderr.InsertExistError } } return err } func sortCheck(sort interface{}) error { if !enableSortMCheck { return nil } if sort == nil { return nil } typ := reflect.TypeOf(sort) var m bson.M switch typ { case bsonMType: //log.Warn("mongo sort use bson.M use bson.D instead", log.Any("sort", sort)) m, _ = sort.(bson.M) case bsonMPtrType: //log.Warn("mongo sort use *bson.M use bson.D instead", log.Any("sort", sort)) pm, _ := sort.(*bson.M) m = *pm case bsonDType, bsonDPtrType, bsonEType, bsonEPtrType: return nil default: log.Warn("sort use unknown sort type please check", log.Any("sort", sort), log.Any("typ", typ)) return errors.New("mongo error sort type") } if len(m) > 1 { return errors.New("mongo error sort, use bson.M and len(sort) > 1") } return nil } // CreateIndex 创建数据索引.options 在index里面创建 func (u *MongoTool) CreateIndex(models []mongo.IndexModel) ([]string, error) { //return nil, nil return u.coll.Indexes().CreateMany(u.ctx, models) } // DropIndex 删除数据索引.options 在index里面创建 func (u *MongoTool) DropIndex(indexname string) error { if _, err := u.coll.Indexes().DropOne(u.ctx, indexname); err != nil { log.Error(fmt.Sprintf("drop indexes error %+v", err)) return err } return nil } // DropIndexIfExists 删除指定索引;索引或所在集合尚未创建时均视为成功。 func (u *MongoTool) DropIndexIfExists(indexname string) error { if _, err := u.coll.Indexes().DropOne(u.ctx, indexname); err != nil { var commandErr mongo.CommandError // 27=IndexNotFound(索引不存在)、26=NamespaceNotFound(集合/库尚未创建):目标索引本就不存在,视为成功。 if errors.As(err, &commandErr) && (commandErr.Code == 27 || commandErr.Code == 26) { return nil } log.Error(fmt.Sprintf("drop indexes error %+v", err)) return err } return nil } // InsertOne 插入单条信息 func (u *MongoTool) InsertOne(document interface{}, opts ...*options.InsertOneOptions) (*mongo.InsertOneResult, error) { insertResult, err := u.coll.InsertOne(u.ctx, document, opts...) return insertResult, handleDbError(err) } // InsertMany 批量插入信息 func (u *MongoTool) InsertMany(documents interface{}, opts ...*options.InsertManyOptions) (*mongo.InsertManyResult, error) { if err := validInterfaceSlice(documents); err != nil { return nil, err } res := setTimeForSlice(documents) return u.coll.InsertMany(u.ctx, res, opts...) } // Find 查询多条数据 func (u *MongoTool) Find(model interface{}, filter bson.M, opts ...*options.FindOptions) error { if err := validInterfaceSlice(model); err != nil { return err } cur, err := u.FindCursor(filter, opts...) if err != nil { return err } return handleDbError(cur.All(u.ctx, model)) } // FindCursor 查询多条数据并返回游标。 // Cursor 不是并发安全的,调用方必须在完成或失败后关闭它。 func (u *MongoTool) FindCursor(filter bson.M, opts ...*options.FindOptions) (*mongo.Cursor, error) { for _, opt := range opts { if opt == nil { continue } if opt.Limit != nil && *opt.Limit > 1000 { fmt.Println("limit beyond 1000 ==================>", *opt.Limit) } if opt.Sort != nil { if err := sortCheck(opt.Sort); err != nil { return nil, err } } } return u.coll.Find(u.ctx, filter, opts...) } // FindOne 单条查询 func (u *MongoTool) FindOne(model interface{}, filter bson.M, opts ...*options.FindOneOptions) error { for _, opt := range opts { if opt.Sort != nil { if err := sortCheck(opt.Sort); err != nil { return err } } } return handleDbError(u.coll.FindOne(u.ctx, filter, opts...).Decode(model)) } // FindOneAndDelete 单条查询并删除 func (u *MongoTool) FindOneAndDelete(model interface{}, filter bson.M, opts ...*options.FindOneAndDeleteOptions) error { for _, opt := range opts { if opt.Sort != nil { if err := sortCheck(opt.Sort); err != nil { return err } } } return handleDbError(u.coll.FindOneAndDelete(u.ctx, filter, opts...).Decode(model)) } // FindOneAndReplace 单条查询 rd set to Before 表示返回原始数据, set to After 表示返回替换后的数据 func (u *MongoTool) FindOneAndReplace(model interface{}, filter bson.M, replacement bson.M, opts ...*options.FindOneAndReplaceOptions) error { for _, opt := range opts { if opt.Sort != nil { if err := sortCheck(opt.Sort); err != nil { return err } } } return handleDbError(u.coll.FindOneAndReplace(u.ctx, filter, replacement, opts...).Decode(model)) } // FindOneAndUpdate 单条查询 rd set to Before 表示返回原始数据, set to After 表示返回更新后的数据 默认为返回更新后的数据 func (u *MongoTool) FindOneAndUpdate(model interface{}, filter bson.M, update bson.M, opts ...*options.FindOneAndUpdateOptions) error { for _, opt := range opts { if opt.Sort != nil { if err := sortCheck(opt.Sort); err != nil { return err } } } 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}) } return handleDbError(u.coll.FindOneAndUpdate(u.ctx, filter, update, opts...).Decode(model)) } // FindOneAndUpsert 单条查询 匹配到数据更新,未匹配到数据则upsert func (u *MongoTool) FindOneAndUpsert(model interface{}, filter bson.M, update bson.M, opts ...*options.FindOneAndUpdateOptions) error { upsert := true var beforeOrAfter options.ReturnDocument for _, opt := range opts { if opt.Sort != nil { if err := sortCheck(opt.Sort); err != nil { return err } } } setReturn := false for _, opt := range opts { if opt != nil && opt.ReturnDocument != nil { beforeOrAfter = *opt.ReturnDocument setReturn = true break } } if !setReturn { beforeOrAfter = options.After } opts = append(opts, &options.FindOneAndUpdateOptions{ReturnDocument: &beforeOrAfter, Upsert: &upsert}) return handleDbError(u.coll.FindOneAndUpdate(u.ctx, filter, update, opts...).Decode(model)) } func (u *MongoTool) FindOneAndUpdateReturnTiny(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 := handleDbError(u.coll.FindOneAndUpdate(u.ctx, query, update, opts...).Decode(bind)) return result } // FindOneByID 通过id查找一条数据 func (u *MongoTool) FindOneByID(model interface{}, id primitive.ObjectID, opts ...*options.FindOneOptions) error { for _, opt := range opts { if opt.Sort != nil { if err := sortCheck(opt.Sort); err != nil { return err } } } return handleDbError(u.coll.FindOne(u.ctx, bson.M{"_id": id}, opts...).Decode(model)) } // Aggregate 聚合查找数据 func (u *MongoTool) Aggregate(model interface{}, pipeline []bson.M, opts ...*options.AggregateOptions) error { if err := validInterfaceSlice(model); err != nil { return err } cur, err := u.coll.Aggregate(u.ctx, pipeline, opts...) if err != nil { return err } return handleDbError(cur.All(u.ctx, model)) } // AggregateDecode 聚合.Decode func (u *MongoTool) AggregateDecode(model interface{}, pipeline []bson.M, opts ...*options.AggregateOptions) error { cur, err := u.coll.Aggregate(u.ctx, pipeline, opts...) if err != nil { return err } if cur.Next(u.ctx) { return handleDbError(cur.Decode(model)) } return nil } // Distinct 去重查询 func (u *MongoTool) Distinct(fieldName string, filter bson.M, opts ...*options.DistinctOptions) ([]interface{}, error) { return u.coll.Distinct(u.ctx, fieldName, filter, opts...) } // DeleteOne 删除一条数据 func (u *MongoTool) DeleteOne(filter bson.M, opt ...*options.DeleteOptions) (*mongo.DeleteResult, error) { return u.coll.DeleteOne(u.ctx, filter, opt...) } // DeleteMany 删除多条数据 func (u *MongoTool) DeleteMany(filter bson.M, opt ...*options.DeleteOptions) (*mongo.DeleteResult, error) { return u.coll.DeleteMany(u.ctx, filter, opt...) } // DeleteById 根据ID删除数据单条数据 func (u *MongoTool) DeleteById(id primitive.ObjectID) (*mongo.DeleteResult, error) { return u.coll.DeleteOne(u.ctx, bson.M{"_id": id}) } // UpdateOne 更新单条数据 func (u *MongoTool) UpdateOne(filter bson.M, update interface{}) (*mongo.UpdateResult, error) { return u.coll.UpdateOne(u.ctx, filter, update) } // UpdateMany 修改多条数据 func (u *MongoTool) UpdateMany(filter bson.M, update bson.M) (*mongo.UpdateResult, error) { return u.coll.UpdateMany(u.ctx, filter, update) } // UpsertMany 或者修改或者插入多条数据 func (u *MongoTool) UpsertMany(filter bson.M, update bson.M) (*mongo.UpdateResult, error) { return u.coll.UpdateMany(u.ctx, filter, update, options.Update().SetUpsert(true)) } // UpsertOne 或者修改或者插入一条数据 func (u *MongoTool) UpsertOne(filter bson.M, update bson.M) (*mongo.UpdateResult, error) { return u.coll.UpdateOne(u.ctx, filter, update, options.Update().SetUpsert(true)) } // UpdateOneForSet 修改一条数据 【根据修改数据中集合类型字段】 func (u *MongoTool) UpdateOneForSet(filter bson.M, update bson.D) (*mongo.UpdateResult, error) { return u.coll.UpdateOne(u.ctx, filter, update) } // Count 获取数量 func (u *MongoTool) Count(filter interface{}, opts ...*options.CountOptions) (int64, error) { if reflect.TypeOf(filter).Kind() == reflect.Slice { cur, err := u.coll.Aggregate(u.ctx, filter) if err != nil { return 0, err } var count int64 = 0 for cur.Next(context.TODO()) { count++ } return count, nil } return u.coll.CountDocuments(u.ctx, filter, opts...) } // EstimateCount 获取数量 func (u *MongoTool) EstimateCount(opts ...*options.EstimatedDocumentCountOptions) (int64, error) { return u.coll.EstimatedDocumentCount(u.ctx, opts...) } // Bulk Bulk func (u *MongoTool) Bulk(models []mongo.WriteModel, opts ...*options.BulkWriteOptions) (*mongo.BulkWriteResult, error) { return u.coll.BulkWrite(u.ctx, models, opts...) } // Exists 是否存在数据 func (u *MongoTool) Exists(filter interface{}, opts ...*options.FindOneOptions) (bool, error) { var limit int64 = 1 lo := &options.CountOptions{Limit: &limit} n, err := u.Count(filter, lo) if err != nil { return false, err } return n > 0, nil } func setTimeForSlice(docs interface{}) []interface{} { arr := reflect.ValueOf(docs) if arr.Kind() == reflect.Ptr { arr = reflect.ValueOf(docs).Elem() } result := make([]interface{}, arr.Len()) for i := 0; i < arr.Len(); i++ { ele := arr.Index(i) 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[i] = ele.Interface() } return result } func validInterfaceSlice(bind interface{}) error { t := reflect.TypeOf(bind) k := t.Kind() if t.Kind() == reflect.Ptr { k = t.Elem().Kind() } if k != reflect.Slice { return stderr.MustSliceOrSlicePtr } return nil }