web_hgtv_server/pub/mongos/base_mod.go

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
}