go_gameshield/pkg/mongo/mongo.go

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
}