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 } counts := make(map[string]int64) for day := startDay; !day.After(endDay); day = day.AddDate(0, 0, 1) { date := day.Format("2006-01-02") value, err := c.Get(ctx, AITranslationSuccessDailyCounterKey(day)) if err != nil { return nil, err } counts[date] = value } 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 "" } }