64 lines
1.7 KiB
Go
64 lines
1.7 KiB
Go
package data
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"strings"
|
|
"time"
|
|
|
|
"kra/internal/config"
|
|
|
|
"go.mongodb.org/mongo-driver/mongo"
|
|
"go.mongodb.org/mongo-driver/mongo/options"
|
|
)
|
|
|
|
func mongoURI(config *config.Mongo) string {
|
|
hosts := make([]string, 0, len(config.Hosts))
|
|
for _, host := range config.Hosts {
|
|
if host != nil && host.Host != "" && host.Port != "" {
|
|
hosts = append(hosts, host.Host+":"+host.Port)
|
|
}
|
|
}
|
|
uri := "mongodb://" + strings.Join(hosts, ",") + "/" + config.Database
|
|
if config.Options != "" {
|
|
uri += "?" + config.Options
|
|
}
|
|
return uri
|
|
}
|
|
|
|
func openMongo(config *config.Mongo, enabled bool) (*mongo.Client, error) {
|
|
if !enabled {
|
|
return nil, nil
|
|
}
|
|
if config == nil || len(config.Hosts) == 0 {
|
|
return nil, fmt.Errorf("mongo hosts are required")
|
|
}
|
|
clientOptions := options.Client().ApplyURI(mongoURI(config))
|
|
if config.Username != "" && config.Password != "" {
|
|
clientOptions.SetAuth(options.Credential{Username: config.Username, Password: config.Password, AuthSource: config.AuthSource})
|
|
}
|
|
if config.MinPoolSize > 0 {
|
|
clientOptions.SetMinPoolSize(config.MinPoolSize)
|
|
}
|
|
if config.MaxPoolSize > 0 {
|
|
clientOptions.SetMaxPoolSize(config.MaxPoolSize)
|
|
}
|
|
if config.ConnectTimeoutMs > 0 {
|
|
clientOptions.SetConnectTimeout(time.Duration(config.ConnectTimeoutMs) * time.Millisecond)
|
|
}
|
|
if config.SocketTimeoutMs > 0 {
|
|
clientOptions.SetSocketTimeout(time.Duration(config.SocketTimeoutMs) * time.Millisecond)
|
|
}
|
|
ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
|
|
defer cancel()
|
|
client, err := mongo.Connect(ctx, clientOptions)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
if err = client.Ping(ctx, nil); err != nil {
|
|
_ = client.Disconnect(context.Background())
|
|
return nil, err
|
|
}
|
|
return client, nil
|
|
}
|