tm_back/pkg/mongo/counter.go

package mongo

import (
	"context"
	"errors"
	"fmt"
	"time"

	"go.mongodb.org/mongo-driver/v2/bson"
	mongodriver "go.mongodb.org/mongo-driver/v2/mongo"
	"go.mongodb.org/mongo-driver/v2/mongo/options"
)

const (
	metricsCountersCollection                   = "metrics_counters"
	AITranslationSuccessCounterKey              = "ai_translation_successful_requests"
	AITranslationSuccessCounterKeyDailyJob      = "ai_translation_successful_requests_daily_job"
	AITranslationSuccessCounterKeyManualTrigger = "ai_translation_successful_requests_manual_trigger"
)

// Counter provides atomic increment/read operations for named metrics stored in MongoDB.
type Counter struct {
	mongo *ConnectionManager
}

// NewCounter creates a counter backed by the metrics_counters collection.
func NewCounter(mongo *ConnectionManager) *Counter {
	return &Counter{mongo: mongo}
}

// Increment atomically increments the counter identified by key and returns the new value.
func (c *Counter) Increment(ctx context.Context, key string) (int64, error) {
	if key == "" {
		return 0, fmt.Errorf("mongo counter: key is required")
	}

	var res struct {
		Count int64 `bson:"count"`
	}
	err := c.mongo.GetCollection(metricsCountersCollection).FindOneAndUpdate(
		ctx,
		bson.M{"_id": key},
		bson.M{
			"$inc":         bson.M{"count": 1},
			"$set":         bson.M{"updated_at": time.Now().Unix()},
			"$setOnInsert": bson.M{"_id": key},
		},
		options.FindOneAndUpdate().SetUpsert(true).SetReturnDocument(options.After),
	).Decode(&res)
	if err != nil {
		return 0, fmt.Errorf("mongo counter increment %q: %w", key, err)
	}

	return res.Count, nil
}

// Get returns the current value of the counter identified by key (0 when missing).
func (c *Counter) Get(ctx context.Context, key string) (int64, error) {
	if key == "" {
		return 0, fmt.Errorf("mongo counter: key is required")
	}

	var doc struct {
		Count int64 `bson:"count"`
	}
	err := c.mongo.GetCollection(metricsCountersCollection).FindOne(ctx, bson.M{"_id": key}).Decode(&doc)
	if err != nil {
		if errors.Is(err, mongodriver.ErrNoDocuments) {
			return 0, nil
		}
		return 0, fmt.Errorf("mongo counter get %q: %w", key, err)
	}

	return doc.Count, nil
}

// AITranslationSuccessCallback returns a callback that increments successful AI translation counters.
func AITranslationSuccessCallback(counter *Counter) func(ctx context.Context, source string) error {
	return func(ctx context.Context, source string) error {
		if _, err := counter.Increment(ctx, AITranslationSuccessCounterKey); err != nil {
			return err
		}
		if _, err := counter.Increment(ctx, AITranslationSuccessDailyCounterKey(time.Now().UTC())); err != nil {
			return err
		}
		if key := AITranslationSuccessCounterKeyForSource(source); key != "" {
			_, err := counter.Increment(ctx, key)
			return err
		}
		return nil
	}
}

// AITranslationSuccessDailyCounterKey returns the metrics key for successful translations on a given UTC day.
func AITranslationSuccessDailyCounterKey(day time.Time) string {
	return AITranslationSuccessCounterKey + "_day:" + day.UTC().Format("2006-01-02")
}

// GetDailyCounts returns counter values for each UTC day from start through end inclusive.
func (c *Counter) GetDailyCounts(ctx context.Context, startDay, endDay time.Time) (map[string]int64, error) {
	startDay = time.Date(startDay.Year(), startDay.Month(), startDay.Day(), 0, 0, 0, 0, time.UTC)
	endDay = time.Date(endDay.Year(), endDay.Month(), endDay.Day(), 0, 0, 0, 0, time.UTC)
	if endDay.Before(startDay) {
		return map[string]int64{}, nil
	}

	keys := make([]string, 0)
	dates := make([]string, 0)
	for day := startDay; !day.After(endDay); day = day.AddDate(0, 0, 1) {
		dates = append(dates, day.Format("2006-01-02"))
		keys = append(keys, AITranslationSuccessDailyCounterKey(day))
	}

	counts := make(map[string]int64, len(dates))
	for _, date := range dates {
		counts[date] = 0
	}

	cursor, err := c.mongo.GetCollection(metricsCountersCollection).Find(ctx, bson.M{
		"_id": bson.M{"$in": keys},
	})
	if err != nil {
		return nil, fmt.Errorf("mongo counter get daily counts: %w", err)
	}
	defer cursor.Close(ctx)

	keyToDate := make(map[string]string, len(keys))
	for i, key := range keys {
		keyToDate[key] = dates[i]
	}

	for cursor.Next(ctx) {
		var doc struct {
			ID    string `bson:"_id"`
			Count int64  `bson:"count"`
		}
		if err := cursor.Decode(&doc); err != nil {
			return nil, fmt.Errorf("mongo counter decode daily count: %w", err)
		}
		date, ok := keyToDate[doc.ID]
		if !ok {
			continue
		}
		counts[date] = doc.Count
	}
	if err := cursor.Err(); err != nil {
		return nil, fmt.Errorf("mongo counter iterate daily counts: %w", err)
	}

	return counts, nil
}

// AITranslationSuccessCounterKeyForSource maps a request source to its metrics counter key.
func AITranslationSuccessCounterKeyForSource(source string) string {
	switch source {
	case "daily_job":
		return AITranslationSuccessCounterKeyDailyJob
	case "manual_trigger":
		return AITranslationSuccessCounterKeyManualTrigger
	default:
		return ""
	}
}