|
@@ -0,0 +1,79 @@
|
|
|
|
+package mongo
|
|
|
|
+
|
|
|
|
+import (
|
|
|
|
+ "context"
|
|
|
|
+ "fmt"
|
|
|
|
+ "time"
|
|
|
|
+
|
|
|
|
+ "go.mongodb.org/mongo-driver/mongo"
|
|
|
|
+ "go.mongodb.org/mongo-driver/mongo/options"
|
|
|
|
+ "go.mongodb.org/mongo-driver/mongo/readpref"
|
|
|
|
+)
|
|
|
|
+
|
|
|
|
+// Config MongoDB配置
|
|
|
|
+type Config struct {
|
|
|
|
+ URI string `mapstructure:"uri"`
|
|
|
|
+ Database string `mapstructure:"database"`
|
|
|
|
+ Timeout time.Duration `mapstructure:"timeout"`
|
|
|
|
+ MaxPoolSize uint64 `mapstructure:"max_pool_size"`
|
|
|
|
+}
|
|
|
|
+
|
|
|
|
+// MongoDB连接管理器
|
|
|
|
+type MongoDB struct {
|
|
|
|
+ config *Config
|
|
|
|
+ client *mongo.Client
|
|
|
|
+ database *mongo.Database
|
|
|
|
+}
|
|
|
|
+
|
|
|
|
+// New 创建新的MongoDB客户端
|
|
|
|
+func New(config *Config) (*MongoDB, error) {
|
|
|
|
+ ctx, cancel := context.WithTimeout(context.Background(), config.Timeout)
|
|
|
|
+ defer cancel()
|
|
|
|
+
|
|
|
|
+ // 创建MongoDB客户端选项
|
|
|
|
+ clientOptions := options.Client().
|
|
|
|
+ ApplyURI(config.URI).
|
|
|
|
+ SetMaxPoolSize(config.MaxPoolSize)
|
|
|
|
+
|
|
|
|
+ // 连接到MongoDB
|
|
|
|
+ client, err := mongo.Connect(ctx, clientOptions)
|
|
|
|
+ if err != nil {
|
|
|
|
+ return nil, fmt.Errorf("连接MongoDB失败: %w", err)
|
|
|
|
+ }
|
|
|
|
+
|
|
|
|
+ // 验证连接
|
|
|
|
+ if err := client.Ping(ctx, readpref.Primary()); err != nil {
|
|
|
|
+ return nil, fmt.Errorf("MongoDB连接测试失败: %w", err)
|
|
|
|
+ }
|
|
|
|
+
|
|
|
|
+ // 获取数据库
|
|
|
|
+ database := client.Database(config.Database)
|
|
|
|
+
|
|
|
|
+ return &MongoDB{
|
|
|
|
+ config: config,
|
|
|
|
+ client: client,
|
|
|
|
+ database: database,
|
|
|
|
+ }, nil
|
|
|
|
+}
|
|
|
|
+
|
|
|
|
+// Close 关闭MongoDB连接
|
|
|
|
+func (m *MongoDB) Close() error {
|
|
|
|
+ ctx, cancel := context.WithTimeout(context.Background(), m.config.Timeout)
|
|
|
|
+ defer cancel()
|
|
|
|
+ return m.client.Disconnect(ctx)
|
|
|
|
+}
|
|
|
|
+
|
|
|
|
+// GetDatabase 获取数据库实例
|
|
|
|
+func (m *MongoDB) GetDatabase() *mongo.Database {
|
|
|
|
+ return m.database
|
|
|
|
+}
|
|
|
|
+
|
|
|
|
+// GetCollection 获取集合实例
|
|
|
|
+func (m *MongoDB) GetCollection(name string) *mongo.Collection {
|
|
|
|
+ return m.database.Collection(name)
|
|
|
|
+}
|
|
|
|
+
|
|
|
|
+// Client 获取原始MongoDB客户端
|
|
|
|
+func (m *MongoDB) Client() *mongo.Client {
|
|
|
|
+ return m.client
|
|
|
|
+}
|