logmower-shipper/cmd/mongo.go

package logmower

import (
	"context"
	"fmt"
	"log"
	"net/url"
	"time"

	ms "git.k-space.ee/k-space/logmower-shipper/pkg/mongoStruct"
	prom "github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promauto"
	mongoEvent "go.mongodb.org/mongo-driver/event"
	"go.mongodb.org/mongo-driver/mongo"
	mongoOpt "go.mongodb.org/mongo-driver/mongo/options"
)

var (
	promDbHeartbeat = promauto.NewHistogramVec(prom.HistogramOpts{
		Namespace: PrometheusPrefix,
		Subsystem: "database",
		Name:      "heartbeat_time",
		Help:      "Time in seconds for succeeded heartbeat, or 0 on failure",
		Buckets:   []float64{0.1, 0.2, 0.5, 1, 5, 10, 50},
	}, []string{"connection_id"})

	promDbCmd = promauto.NewHistogramVec(prom.HistogramOpts{
		Namespace: PrometheusPrefix,
		Subsystem: "database",
		Name:      "operation_latency", // "command_time",
		Help:      "Time in seconds of commands",
		Buckets:   []float64{0.1, 0.2, 0.5, 1, 5, 10, 50},
	}, []string{"connection_id", "command_name"})

	promDbCmdErr = promauto.NewCounterVec(prom.CounterOpts{
		Namespace: PrometheusPrefix,
		Subsystem: "database",
		Name:      "errors",
		Help:      "Failed commands (also reflected elsewhere)",
	}, []string{"connection_id", "command_name"})
)

func mongoMonitoredClientOptions() *mongoOpt.ClientOptions {
	return mongoOpt.Client().
		SetServerMonitor(&mongoEvent.ServerMonitor{
			ServerHeartbeatSucceeded: func(ev *mongoEvent.ServerHeartbeatSucceededEvent) {
				promDbHeartbeat.WithLabelValues(ev.ConnectionID).Observe(time.Duration(ev.DurationNanos).Seconds())
			},
			ServerHeartbeatFailed: func(ev *mongoEvent.ServerHeartbeatFailedEvent) {
				promDbHeartbeat.WithLabelValues(ev.ConnectionID).Observe(0)
				log.Printf("database heartbeat failed on connection %q: %e", ev.ConnectionID, ev.Failure)
			},
		}).
		SetMonitor(&mongoEvent.CommandMonitor{
			Succeeded: func(_ context.Context, ev *mongoEvent.CommandSucceededEvent) {
				promDbCmd.WithLabelValues(ev.ConnectionID, ev.CommandName).Observe(time.Duration(ev.DurationNanos).Seconds())
			},
			Failed: func(_ context.Context, ev *mongoEvent.CommandFailedEvent) {
				promDbCmd.WithLabelValues(ev.ConnectionID, ev.CommandName).Observe(time.Duration(ev.DurationNanos).Seconds())

				promDbCmdErr.WithLabelValues(ev.ConnectionID, ev.CommandName).Add(1)
			},
		})
}

func initDatabase(ctx context.Context, uri string) (*mongo.Collection, error) {
	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 := mongoMonitoredClientOptions().ApplyURI(uri)

	dbClient, err := mongo.Connect(mongoTimeoutCtx(ctx))
	if err != nil {
		return nil, fmt.Errorf("connecting to %q: %w", dbOpt.GetURI(), err)
	}

	col := dbClient.Database(uriParsed.Path).Collection("logs")

	if err := ms.InitializeIndexes(mongoTimeoutCtx(ctx), col); err != nil {
		return nil, fmt.Errorf("initializing indexes: %w", err)
	}

	return col, nil
}