logmower-shipper/pkg/mongo/mongo.go

73 lines
1.9 KiB
Go

package mongo
import (
"context"
"fmt"
"net/url"
"time"
"go.mongodb.org/mongo-driver/mongo"
mongoOpt "go.mongodb.org/mongo-driver/mongo/options"
)
const CommandTimeout = 10 * time.Second
func GlobalTimeout(ctx context.Context) context.Context {
ctx, _ = context.WithTimeout(ctx, CommandTimeout) //nolint:lostcancel (cancelled by mongo, should be bug on them //TODO)
return ctx
}
type InitializeOptions struct {
MaxPoolSize uint64
CapSizeBytes int64
ExpireAfterSeconds int64
}
func Initialize(ctx context.Context, uri string, opt *InitializeOptions) (*mongo.Collection, error) {
collectionName := "logs"
if opt == nil {
opt = &InitializeOptions{}
}
uriParsed, err := url.ParseRequestURI(uri)
if err != nil {
return nil, fmt.Errorf("parsing URI for database name: %w", err)
}
uriParsed.Path = uriParsed.Path[1:] // remove leading slash
if uriParsed.Path == "" {
return nil, fmt.Errorf("URI must include database name (as database to authenticate against)")
}
dbOpt := attachMetrics(mongoOpt.Client()).ApplyURI(uri).SetMaxPoolSize(opt.MaxPoolSize)
dbClient, err := mongo.Connect(GlobalTimeout(ctx), dbOpt)
if err != nil {
return nil, fmt.Errorf("connecting to %q: %w", dbOpt.GetURI(), err)
}
if err := dbClient.Ping(GlobalTimeout(ctx), nil); err != nil {
return nil, fmt.Errorf("first ping to database: %w", err)
}
db := dbClient.Database(uriParsed.Path)
capped := opt.CapSizeBytes > 0
if err := db.CreateCollection(GlobalTimeout(ctx), collectionName, &mongoOpt.CreateCollectionOptions{
Capped: &capped,
SizeInBytes: &opt.CapSizeBytes,
ExpireAfterSeconds: &opt.ExpireAfterSeconds,
}); err != nil {
return nil, fmt.Errorf("initializing collection")
}
col := db.Collection(collectionName)
if err := InitializeIndexes(GlobalTimeout(ctx), col); err != nil {
return nil, fmt.Errorf("initializing indexes: %w", err)
}
return col, nil
}