Files
tm_back/pkg/mongo/connection.go
T
2026-04-21 19:07:51 +03:30

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
}