Initial commit

2026-05-20 16:13:58 +08:00
commit ceffbe382e
500 changed files with 157469 additions and 0 deletions
@@ -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
+}