284 lines
8.1 KiB
Go
284 lines
8.1 KiB
Go
package mongo
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"time"
|
|
|
|
"go.mongodb.org/mongo-driver/v2/bson"
|
|
"go.mongodb.org/mongo-driver/v2/mongo"
|
|
"go.mongodb.org/mongo-driver/v2/mongo/options"
|
|
"go.mongodb.org/mongo-driver/v2/mongo/readpref"
|
|
)
|
|
|
|
// ConnectionConfig holds MongoDB connection configuration
|
|
type ConnectionConfig struct {
|
|
URI string `json:"uri"`
|
|
Database string `json:"database"`
|
|
MaxPoolSize uint64 `json:"maxPoolSize"`
|
|
MinPoolSize uint64 `json:"minPoolSize"`
|
|
MaxConnIdleTime time.Duration `json:"maxConnIdleTime"`
|
|
ConnectTimeout time.Duration `json:"connectTimeout"`
|
|
ServerSelectionTimeout time.Duration `json:"serverSelectionTimeout"`
|
|
}
|
|
|
|
// ConnectionManager manages MongoDB connections and collections
|
|
type ConnectionManager struct {
|
|
client *mongo.Client
|
|
database *mongo.Database
|
|
logger Logger
|
|
config ConnectionConfig
|
|
}
|
|
|
|
// NewConnectionManager creates a new MongoDB connection manager
|
|
func NewConnectionManager(config ConnectionConfig, logger Logger) (*ConnectionManager, error) {
|
|
// Set default values
|
|
if config.MaxPoolSize == 0 {
|
|
config.MaxPoolSize = 100
|
|
}
|
|
if config.MinPoolSize == 0 {
|
|
config.MinPoolSize = 5
|
|
}
|
|
if config.MaxConnIdleTime == 0 {
|
|
config.MaxConnIdleTime = 30 * time.Minute
|
|
}
|
|
if config.ConnectTimeout == 0 {
|
|
config.ConnectTimeout = 10 * time.Second
|
|
}
|
|
if config.ServerSelectionTimeout == 0 {
|
|
config.ServerSelectionTimeout = 5 * time.Second
|
|
}
|
|
|
|
// Create client options
|
|
clientOptions := options.Client().
|
|
ApplyURI(config.URI).
|
|
SetMaxPoolSize(config.MaxPoolSize).
|
|
SetMinPoolSize(config.MinPoolSize).
|
|
SetMaxConnIdleTime(config.MaxConnIdleTime).
|
|
SetConnectTimeout(config.ConnectTimeout).
|
|
SetServerSelectionTimeout(config.ServerSelectionTimeout)
|
|
|
|
// Connect to MongoDB
|
|
ctx, cancel := context.WithTimeout(context.Background(), config.ConnectTimeout)
|
|
defer cancel()
|
|
|
|
client, err := mongo.Connect(clientOptions)
|
|
if err != nil {
|
|
logger.Error("Failed to connect to MongoDB", map[string]interface{}{
|
|
"uri": config.URI,
|
|
"error": err.Error(),
|
|
})
|
|
return nil, fmt.Errorf("failed to connect to MongoDB: %w", err)
|
|
}
|
|
|
|
// Ping the database
|
|
if err = client.Ping(ctx, readpref.Primary()); err != nil {
|
|
logger.Error("Failed to ping MongoDB", map[string]interface{}{
|
|
"error": err.Error(),
|
|
})
|
|
return nil, fmt.Errorf("failed to ping MongoDB: %w", err)
|
|
}
|
|
|
|
database := client.Database(config.Database)
|
|
|
|
logger.Info("Successfully connected to MongoDB", map[string]interface{}{
|
|
"database": config.Database,
|
|
"uri": config.URI,
|
|
})
|
|
|
|
return &ConnectionManager{
|
|
client: client,
|
|
database: database,
|
|
logger: logger,
|
|
config: config,
|
|
}, nil
|
|
}
|
|
|
|
// GetCollection returns a MongoDB collection
|
|
func (cm *ConnectionManager) GetCollection(name string) *mongo.Collection {
|
|
return cm.database.Collection(name)
|
|
}
|
|
|
|
// CreateRepository creates a new repository for a given collection
|
|
func (cm *ConnectionManager) CreateRepository(collectionName string) *mongo.Collection {
|
|
return cm.GetCollection(collectionName)
|
|
}
|
|
|
|
// CreateIndexes creates indexes for a collection
|
|
func (cm *ConnectionManager) CreateIndexes(collectionName string, indexes []Index) error {
|
|
collection := cm.GetCollection(collectionName)
|
|
|
|
for _, index := range indexes {
|
|
indexModel := mongo.IndexModel{
|
|
Keys: index.Keys,
|
|
Options: index.Options,
|
|
}
|
|
|
|
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
|
|
defer cancel()
|
|
|
|
_, err := collection.Indexes().CreateOne(ctx, indexModel)
|
|
if err != nil {
|
|
cm.logger.Error("Failed to create index", map[string]interface{}{
|
|
"collection": collectionName,
|
|
"index": index.Name,
|
|
"error": err.Error(),
|
|
})
|
|
return fmt.Errorf("failed to create index %s: %w", index.Name, err)
|
|
}
|
|
|
|
cm.logger.Info("Index created successfully", map[string]interface{}{
|
|
"collection": collectionName,
|
|
"index": index.Name,
|
|
})
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
// DropIndexes drops indexes from a collection
|
|
func (cm *ConnectionManager) DropIndexes(collectionName string, indexNames []string) error {
|
|
collection := cm.GetCollection(collectionName)
|
|
|
|
for _, indexName := range indexNames {
|
|
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
|
|
defer cancel()
|
|
|
|
err := collection.Indexes().DropOne(ctx, indexName)
|
|
if err != nil {
|
|
cm.logger.Error("Failed to drop index", map[string]interface{}{
|
|
"collection": collectionName,
|
|
"index": indexName,
|
|
"error": err.Error(),
|
|
})
|
|
return fmt.Errorf("failed to drop index %s: %w", indexName, err)
|
|
}
|
|
|
|
cm.logger.Info("Index dropped successfully", map[string]interface{}{
|
|
"collection": collectionName,
|
|
"index": indexName,
|
|
})
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
// ListIndexes lists all indexes for a collection
|
|
func (cm *ConnectionManager) ListIndexes(collectionName string) ([]bson.M, error) {
|
|
collection := cm.GetCollection(collectionName)
|
|
|
|
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
|
|
defer cancel()
|
|
|
|
cursor, err := collection.Indexes().List(ctx)
|
|
if err != nil {
|
|
cm.logger.Error("Failed to list indexes", map[string]interface{}{
|
|
"collection": collectionName,
|
|
"error": err.Error(),
|
|
})
|
|
return nil, fmt.Errorf("failed to list indexes: %w", err)
|
|
}
|
|
defer cursor.Close(ctx)
|
|
|
|
var indexes []bson.M
|
|
if err = cursor.All(ctx, &indexes); err != nil {
|
|
cm.logger.Error("Failed to decode indexes", map[string]interface{}{
|
|
"error": err.Error(),
|
|
})
|
|
return nil, fmt.Errorf("failed to decode indexes: %w", err)
|
|
}
|
|
|
|
return indexes, nil
|
|
}
|
|
|
|
// Close closes the MongoDB connection
|
|
func (cm *ConnectionManager) Close(ctx context.Context) error {
|
|
if err := cm.client.Disconnect(ctx); err != nil {
|
|
cm.logger.Error("Failed to disconnect from MongoDB", map[string]interface{}{
|
|
"error": err.Error(),
|
|
})
|
|
return fmt.Errorf("failed to disconnect from MongoDB: %w", err)
|
|
}
|
|
|
|
cm.logger.Info("Successfully disconnected from MongoDB", nil)
|
|
return nil
|
|
}
|
|
|
|
// IsConnected checks if MongoDB connection is active
|
|
func (cm *ConnectionManager) IsConnected() bool {
|
|
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
|
|
defer cancel()
|
|
|
|
err := cm.Ping(ctx)
|
|
return err == nil
|
|
}
|
|
|
|
// GetDatabase returns the MongoDB database instance
|
|
func (cm *ConnectionManager) GetDatabase() *mongo.Database {
|
|
return cm.database
|
|
}
|
|
|
|
// GetClient returns the MongoDB client instance
|
|
func (cm *ConnectionManager) GetClient() *mongo.Client {
|
|
return cm.client
|
|
}
|
|
|
|
// Ping pings the MongoDB server
|
|
func (cm *ConnectionManager) Ping(ctx context.Context) error {
|
|
if err := cm.client.Ping(ctx, readpref.Primary()); err != nil {
|
|
cm.logger.Error("Failed to ping MongoDB", map[string]interface{}{
|
|
"error": err.Error(),
|
|
})
|
|
return fmt.Errorf("failed to ping MongoDB: %w", err)
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
// GetStats returns database statistics
|
|
func (cm *ConnectionManager) GetStats(ctx context.Context) (bson.M, error) {
|
|
stats := cm.database.RunCommand(ctx, bson.M{"dbStats": 1})
|
|
if stats.Err() != nil {
|
|
cm.logger.Error("Failed to get database stats", map[string]interface{}{
|
|
"error": stats.Err().Error(),
|
|
})
|
|
return nil, fmt.Errorf("failed to get database stats: %w", stats.Err())
|
|
}
|
|
|
|
var result bson.M
|
|
if err := stats.Decode(&result); err != nil {
|
|
cm.logger.Error("Failed to decode database stats", map[string]interface{}{
|
|
"error": err.Error(),
|
|
})
|
|
return nil, fmt.Errorf("failed to decode database stats: %w", err)
|
|
}
|
|
|
|
return result, nil
|
|
}
|
|
|
|
// GetCollectionStats returns collection statistics
|
|
func (cm *ConnectionManager) GetCollectionStats(ctx context.Context, collectionName string) (bson.M, error) {
|
|
collection := cm.GetCollection(collectionName)
|
|
|
|
stats := collection.Database().RunCommand(ctx, bson.M{
|
|
"collStats": collectionName,
|
|
})
|
|
if stats.Err() != nil {
|
|
cm.logger.Error("Failed to get collection stats", map[string]interface{}{
|
|
"collection": collectionName,
|
|
"error": stats.Err().Error(),
|
|
})
|
|
return nil, fmt.Errorf("failed to get collection stats: %w", stats.Err())
|
|
}
|
|
|
|
var result bson.M
|
|
if err := stats.Decode(&result); err != nil {
|
|
cm.logger.Error("Failed to decode collection stats", map[string]interface{}{
|
|
"error": err.Error(),
|
|
})
|
|
return nil, fmt.Errorf("failed to decode collection stats: %w", err)
|
|
}
|
|
|
|
return result, nil
|
|
}
|