当前位置：首页 > article >正文

goframe开发一个企业网站 MongoDB 完整工具包18

article 2024/11/29 1:53:36

1. MongoDB 工具包完整实现 (mongodb.go)

package mongodb

import (
    "context"
    "fmt"
    "time"
    "github.com/gogf/gf/v2/frame/g"
    "go.mongodb.org/mongo-driver/mongo"
    "go.mongodb.org/mongo-driver/mongo/options"
)

var (
    Client   *mongo.Client
    Database string
)

// Init MongoDB初始化连接
func Init() error {
    ctx := context.Background()
    cfg := g.Cfg()

    // 获取配置参数
    uri := cfg.MustGet(ctx, "mongodb.uri").String()
    Database = cfg.MustGet(ctx, "mongodb.database").String()
    username := cfg.MustGet(ctx, "mongodb.username").String()
    password := cfg.MustGet(ctx, "mongodb.password").String()
    maxPoolSize := cfg.MustGet(ctx, "mongodb.maxPoolSize").Uint64()
    minPoolSize := cfg.MustGet(ctx, "mongodb.minPoolSize").Uint64()
    maxConnIdleTime := cfg.MustGet(ctx, "mongodb.maxConnIdleTime").Int64()

    // 创建连接选项
    clientOptions := options.Client().
        ApplyURI(uri).
        SetAuth(options.Credential{
            Username: username,
            Password: password,
        }).
        SetMaxPoolSize(maxPoolSize).
        SetMinPoolSize(minPoolSize).
        SetMaxConnIdleTime(time.Duration(maxConnIdleTime) * time.Second)

    // 建立连接
    Client, err = mongo.Connect(ctx, clientOptions)
    if err != nil {
        return err
    }

    // 测试连接
    if err := Client.Ping(ctx, nil); err != nil {
        return err
    }

    return nil
}

// IsInitialized 检查是否已初始化
func IsInitialized() bool {
    return Client != nil
}

// GetCollection 获取集合
func GetCollection(collection string) (*mongo.Collection, error) {
    if !IsInitialized() {
        return nil, fmt.Errorf("MongoDB client is not initialized")
    }
    return Client.Database(Database).Collection(collection), nil
}

// Create 创建文档
func Create(collection string, document interface{}) (*mongo.InsertOneResult, error) {
    coll, err := GetCollection(collection)
    if err != nil {
        return nil, err
    }
    return coll.InsertOne(context.Background(), document)
}

// FindOne 查询单个文档
func FindOne(collection string, filter interface{}) (*mongo.SingleResult, error) {
    coll, err := GetCollection(collection)
    if err != nil {
        return nil, err
    }
    return coll.FindOne(context.Background(), filter), nil
}

// Find 查询多个文档（支持分页）
func Find(collection string, filter interface{}, page, pageSize *int64, opts ...*options.FindOptions) (*mongo.Cursor, int64, error) {
    coll, err := GetCollection(collection)
    if err != nil {
        return nil, 0, err
    }

    // 获取总文档数
    total, err := coll.CountDocuments(context.Background(), filter)
    if err != nil {
        return nil, 0, err
    }

    // 计算总页数
    var totalPages int64 = 0
    if page != nil && pageSize != nil && *pageSize > 0 {
        totalPages = (total + *pageSize - 1) / *pageSize
    }

    // 添加分页选项
    if page != nil && pageSize != nil && *page > 0 && *pageSize > 0 {
        skip := (*page - 1) * *pageSize
        opts = append(opts, options.Find().SetSkip(skip).SetLimit(*pageSize))
    }

    cursor, err := coll.Find(context.Background(), filter, opts...)
    return cursor, totalPages, err
}

// Update 更新文档
func Update(collection string, filter interface{}, update interface{}) (*mongo.UpdateResult, error) {
    coll, err := GetCollection(collection)
    if err != nil {
        return nil, err
    }
    return coll.UpdateOne(context.Background(), filter, update)
}

// Delete 删除文档
func Delete(collection string, filter interface{}) (*mongo.DeleteResult, error) {
    coll, err := GetCollection(collection)
    if err != nil {
        return nil, err
    }
    return coll.DeleteOne(context.Background(), filter)
}

// Close 关闭连接
func Close() {
    if Client != nil {
        Client.Disconnect(context.Background())
    }
}

// GetClient 获取MongoDB客户端实例
func GetClient() *mongo.Client {
    return Client
}

2. 功能详解

2.1 初始化相关函数

Init()

功能：初始化 MongoDB 连接
特点：
- 从 GoFrame 配置文件读取配置
- 支持连接池配置
- 支持认证
- 自动测试连接

func Init() error {
    // ... 配置读取 ...
    clientOptions := options.Client().
        ApplyURI(uri).
        SetAuth(options.Credential{
            Username: username,
            Password: password,
        })
    // ... 连接创建 ...
}

IsInitialized()

功能：检查 MongoDB 客户端是否已初始化
使用场景：在执行操作前检查连接状态

2.2 集合操作函数

GetCollection()

功能：获取指定名称的集合
参数：
- collection: 集合名称
返回：
- *mongo.Collection: 集合对象
- error: 错误信息

2.3 CRUD 操作函数

Create()

功能：创建新文档
使用示例：

doc := map[string]interface{}{
    "title": "测试标题",
    "content": "测试内容",
}
result, err := mongodb.Create("collection_name", doc)

FindOne()

功能：查询单个文档
使用示例：

filter := map[string]interface{}{"title": "测试标题"}
result, err := mongodb.FindOne("collection_name", filter)

Find()

功能：查询多个文档（支持分页）
特点：
- 自动计算总页数
- 支持分页查询
- 支持自定义查询选项
使用示例：

page := int64(1)
pageSize := int64(10)
results, totalPages, err := mongodb.Find("collection_name", filter, &page, &pageSize)

Update()

功能：更新文档
使用示例：

filter := map[string]interface{}{"title": "旧标题"}
update := map[string]interface{}{"$set": map[string]interface{}{"title": "新标题"}}
result, err := mongodb.Update("collection_name", filter, update)

Delete()

功能：删除文档
使用示例：

filter := map[string]interface{}{"title": "要删除的标题"}
result, err := mongodb.Delete("collection_name", filter)

3. 使用建议

3.1 错误处理

if err := mongodb.Init(); err != nil {
    g.Log().Fatal(ctx, "MongoDB初始化失败:", err)
}

3.2 连接管理

// 程序启动时初始化
func main() {
    if err := mongodb.Init(); err != nil {
        panic(err)
    }
    defer mongodb.Close()
    // ... 其他代码 ...
}

3.3 分页查询

func GetPagedData(page, pageSize int64) {
    filter := map[string]interface{}{}
    cursor, totalPages, err := mongodb.Find("collection", filter, &page, &pageSize)
    if err != nil {
        // 错误处理
        return
    }
    
    var results []map[string]interface{}
    if err = cursor.All(context.Background(), &results); err != nil {
        // 错误处理
        return
    }
}