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...) }