common/mongo/mongo.go

package mongo

import (
	"context"
	"strings"
	"time"

	"gitee.com/red-future---jilin-g/common/utils"
	"github.com/gogf/gf/v2/errors/gerror"
	"github.com/gogf/gf/v2/frame/g"
	"github.com/gogf/gf/v2/os/glog"
	"github.com/gogf/gf/v2/os/gtime"
	"github.com/gogf/gf/v2/text/gstr"
	"github.com/gogf/gf/v2/util/gconv"
	"go.mongodb.org/mongo-driver/v2/bson"
	"go.mongodb.org/mongo-driver/v2/mongo"
	"go.mongodb.org/mongo-driver/v2/mongo/options"
)

var db = new(mongo.Database)

func init() {
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()
	mongoAddr := g.Cfg().MustGet(context.Background(), "mongo.address").String()
	opt := options.Client().ApplyURI(mongoAddr)
	client, err := mongo.Connect(opt)
	if err != nil {
		glog.Error(ctx, "mongodb连接失败", err)
	}
	// 从连接串中解析数据库名
	dbName := gstr.SubStr(mongoAddr, strings.LastIndex(mongoAddr, "/")+1, len(mongoAddr))
	// 如果连接串带有参数（如 ?retryWrites=true），需要去掉参数部分
	if strings.Contains(dbName, "?") {
		dbName = gstr.SubStr(dbName, 0, strings.Index(dbName, "?"))
	}
	db = client.Database(dbName)
}

// Find 查询多条记录
func Find(ctx context.Context, filter bson.M, result interface{}, collection string, opts ...options.Lister[options.FindOptions]) (err error) {
	if err = utils.ValidStructPtr(result); err != nil {
		return
	}
	filter["isDeleted"] = false
	user, err := utils.GetUserInfo(ctx)
	if err != nil {
		return
	}
	filter["tenantId"] = user.TenantId
	cur, err := db.Collection(collection).Find(ctx, filter, opts...)
	if err != nil {
		return
	}
	err = cur.All(ctx, result)
	return
}

// FindOne 查询1条记录
func FindOne(ctx context.Context, filter bson.M, result interface{}, collection string, opts ...options.Lister[options.FindOneOptions]) (err error) {
	if len(filter) == 0 {
		err = gerror.New("缺少查询条件")
		return
	}
	if err = utils.ValidStructPtr(result); err != nil {
		return
	}
	filter["isDeleted"] = false
	user, err := utils.GetUserInfo(ctx)
	if err != nil {
		return
	}
	filter["tenantId"] = user.TenantId
	cur := db.Collection(collection).FindOne(ctx, filter, opts...)
	err = cur.Decode(result)
	if err == mongo.ErrNoDocuments {
		err = nil
	}
	return
}

// Delete 删除记录
func Delete(ctx context.Context, filter bson.M, collection string, opts ...options.Lister[options.DeleteManyOptions]) (count int64, err error) {
	if len(filter) == 0 {
		err = gerror.New("缺少查询条件")
		return
	}
	user, err := utils.GetUserInfo(ctx)
	if err != nil {
		return
	}
	filter["tenantId"] = user.TenantId
	r, err := db.Collection(collection).DeleteMany(ctx, filter, opts...)
	if err != nil {
		return
	}
	count = r.DeletedCount
	return
}

// Update 修改记录
func Update(ctx context.Context, filter bson.M, update bson.M, collection string, opts ...options.Lister[options.UpdateManyOptions]) (result *mongo.UpdateResult, err error) {
	if len(filter) == 0 {
		err = gerror.New("缺少查询条件")
		return
	}
	filter["isDeleted"] = false
	user, err := utils.GetUserInfo(ctx)
	if err != nil {
		return
	}
	filter["tenantId"] = user.TenantId
	setDoc := update["$set"].(bson.M)
	setDoc["updater"] = user.UserName
	setDoc["updatedAt"] = gtime.Now().Time
	update = bson.M{"$set": setDoc}
	result, err = db.Collection(collection).UpdateMany(ctx, filter, update, opts...)
	if err != nil {
		return
	}
	return
}

// Insert 插入多条记录
func Insert(ctx context.Context, documents []interface{}, collection string, opts ...options.Lister[options.InsertManyOptions]) (ids []interface{}, err error) {
	user, err := utils.GetUserInfo(ctx)
	if err != nil {
		return
	}
	docs := make([]interface{}, 0, len(documents))
	for _, document := range documents {
		doc := gconv.Map(document)
		delete(doc, "id")
		doc["creator"] = user.UserName
		doc["createdAt"] = gtime.Now().Time
		doc["updater"] = user.UserName
		doc["updatedAt"] = gtime.Now().Time
		doc["tenantId"] = user.TenantId
		doc["isDeleted"] = false
		docs = append(docs, doc)
	}
	r, err := db.Collection(collection).InsertMany(ctx, docs, opts...)
	if err != nil {
		return
	}
	ids = r.InsertedIDs
	return
}

// Count  查询总数
func Count(ctx context.Context, filter bson.M, collection string) (count int64, err error) {
	filter["isDeleted"] = false
	// 调用驱动的 CountDocuments，在数据库端执行的
	count, err = db.Collection(collection).CountDocuments(ctx, filter)
	return
}