optimized_go_tools/am_mongo/mongo.go
2024-07-31 09:32:39 +08:00

93 lines
2.3 KiB
Go

package am_mongo
import (
"context"
"fmt"
"go.mongodb.org/mongo-driver/bson"
"go.mongodb.org/mongo-driver/mongo"
"go.mongodb.org/mongo-driver/mongo/options"
)
type MongoDB struct {
Host string
Port string
Username string
Password string
Client *mongo.Client
DB *mongo.Database
}
// 连接到 MongoDB
func (m *MongoDB) Connect(databaseName string) error {
uri := fmt.Sprintf("mongodb://%s:%s@%s:%s", m.Username, m.Password, m.Host, m.Port)
clientOpts := options.Client().ApplyURI(uri)
client, err := mongo.Connect(context.TODO(), clientOpts)
if err != nil {
return err
}
// 检查连接
err = client.Ping(context.TODO(), nil)
if err != nil {
return err
}
m.Client = client
m.DB = client.Database(databaseName)
return nil
}
// 关闭连接
func (m *MongoDB) Close() error {
return m.Client.Disconnect(context.TODO())
}
// 插入一条数据
func (m *MongoDB) Insert(collectionName string, document interface{}) error {
collection := m.DB.Collection(collectionName)
_, err := collection.InsertOne(context.TODO(), document)
return err
}
// 删除数据
func (m *MongoDB) Delete(collectionName string, filter interface{}) error {
collection := m.DB.Collection(collectionName)
_, err := collection.DeleteOne(context.TODO(), filter)
return err
}
// 更新数据
func (m *MongoDB) Update(collectionName string, filter, update interface{}) error {
collection := m.DB.Collection(collectionName)
_, err := collection.UpdateOne(context.TODO(), filter, update)
return err
}
// 查询数据
func (m *MongoDB) Find(collectionName string, filter interface{}, result interface{}) error {
collection := m.DB.Collection(collectionName)
cursor, err := collection.Find(context.TODO(), filter)
if err != nil {
return err
}
defer func(cursor *mongo.Cursor, ctx context.Context) {
_ = cursor.Close(ctx)
}(cursor, context.TODO())
if err = cursor.All(context.TODO(), result); err != nil {
return err
}
return nil
}
// 创建唯一索引
func (m *MongoDB) CreateUniqueIndex(collectionName, fieldName string) error {
collection := m.DB.Collection(collectionName)
indexModel := mongo.IndexModel{
Keys: bson.D{{Key: fieldName, Value: 1}},
Options: options.Index().SetUnique(true),
}
_, err := collection.Indexes().CreateOne(context.TODO(), indexModel)
return err
}