standard/mongodb/tool/mongo.go

package tool

import (
	"context"
	"fmt"

	"go.mongodb.org/mongo-driver/bson"
	"go.mongodb.org/mongo-driver/mongo"
	"go.mongodb.org/mongo-driver/mongo/options"
)

type Config struct {
	Uri      string
	Database string
	Username string
	Password string
}

// MongoRepo 保存 MongoDB 数据库实例的结构体
type MongoRepo struct {
	DB *mongo.Database
}

// InitMongo 初始化 MongoDB 连接，成功后返回 MongoRepo 结构体
func InitMongo(cfg Config) (*MongoRepo, error) {
	// 设置 MongoDB 客户端选项
	clientOptions := options.Client().ApplyURI(cfg.Uri)
	if cfg.Username != "" && cfg.Password != "" {
		clientOptions = clientOptions.SetAuth(options.Credential{
			Username: cfg.Username,
			Password: cfg.Password,
		})
	}

	// 连接到 MongoDB
	client, err := mongo.Connect(context.TODO(), clientOptions)
	if err != nil {
		return nil, err
	}

	// 检查连接
	err = client.Ping(context.TODO(), nil)
	if err != nil {
		// 连接失败时断开连接
		_ = client.Disconnect(context.TODO())
		return nil, err
	}

	// 获取数据库实例
	db := client.Database(cfg.Database)

	return &MongoRepo{
		DB: db,
	}, nil
}

// InsertOne 向指定集合中插入单个文档
func (mr *MongoRepo) InsertOne(ctx context.Context, collectionName string, document interface{}) (*mongo.InsertOneResult, error) {
	collection := mr.DB.Collection(collectionName)
	return collection.InsertOne(ctx, document)
}

// Upsert 执行更新时插入操作
/*
在 MongoDB 里， upsert 是一种特殊的更新行为：
- 若更新操作的查询条件能匹配到文档，就按照指定的更新规则修改这些文档。
- 若查询条件没有匹配到任何文档，就会根据查询条件和更新规则插入一个新的文档。
*/
func (mr *MongoRepo) Upsert(ctx context.Context, collectionName string, filter interface{}, update interface{}) (*mongo.UpdateResult, error) {
	// 获取集合
	collection := mr.DB.Collection(collectionName)
	// 设置更新选项，开启 upsert
	opts := options.Update().SetUpsert(true)
	// 执行更新操作
	return collection.UpdateOne(ctx, filter, update, opts)
}

// NestedQuery 执行嵌套查询，返回匹配的文档游标 如address.city
func (mr *MongoRepo) NestedQuery(ctx context.Context, collectionName string, nestedFilter interface{}, opts ...*options.FindOptions) (*mongo.Cursor, error) {
	collection := mr.DB.Collection(collectionName)
	return collection.Find(ctx, nestedFilter, opts...)
}

// QueryEmbeddedDocument 查询嵌入式文档，返回匹配的文档游标
func (mr *MongoRepo) QueryEmbeddedDocument(ctx context.Context, collectionName string, embeddedFilter interface{}, opts ...*options.FindOptions) (*mongo.Cursor, error) {
	collection := mr.DB.Collection(collectionName)
	return collection.Find(ctx, embeddedFilter, opts...)
}

// QueryArrayContains 查询数组字段包含特定元素的文档，返回匹配的文档游标
func (mr *MongoRepo) QueryArrayContains(ctx context.Context, collectionName string, arrayField string, elements []interface{}, opts ...*options.FindOptions) (*mongo.Cursor, error) {
	filter := bson.M{
		arrayField: bson.M{
			"$in": elements,
		},
	}
	collection := mr.DB.Collection(collectionName)
	return collection.Find(ctx, filter, opts...)
}

// QueryArrayWithComplexFilter 通过复合筛选条件查询数组元素，返回匹配的文档游标
func (mr *MongoRepo) QueryArrayWithComplexFilter(ctx context.Context, collectionName string, arrayField string, complexFilter bson.M, opts ...*options.FindOptions) (*mongo.Cursor, error) {
	filter := bson.M{
		arrayField: bson.M{
			"$elemMatch": complexFilter,
		},
	}
	collection := mr.DB.Collection(collectionName)
	return collection.Find(ctx, filter, opts...)
}

// QueryArrayElementByIndex 按数组索引位置查询元素，返回匹配的文档游标
func (mr *MongoRepo) QueryArrayElementByIndex(ctx context.Context, collectionName string, arrayField string, index int, opts ...*options.FindOptions) (*mongo.Cursor, error) {
	// 构建查询字段，使用点表示法访问数组索引
	field := fmt.Sprintf("%s.%d", arrayField, index)
	// 构建查询条件，这里查询该索引位置有值的文档
	filter := bson.M{
		field: bson.M{
			"$exists": true,
		},
	}
	collection := mr.DB.Collection(collectionName)
	return collection.Find(ctx, filter, opts...)
}