Prometheus Daily Counts Optimized

main.go

package main

import (
	"context"
	"fmt"
	"log"
	"sort"
	"sync"
	"time"

	"github.com/prometheus/client_golang/api"
	v1 "github.com/prometheus/client_golang/api/prometheus/v1"
	"github.com/prometheus/common/model"
	"golang.org/x/sync/errgroup"
)

const usageDateFormat = "2006-01-02"

// DayResult represents the maximum value for a specific day
type DayResult struct {
	Date  time.Time
	Count float64
}

func main() {
	// Create a new Prometheus API client
	client, err := api.NewClient(api.Config{
		Address: "http://localhost:9090",
	})
	if err != nil {
		log.Fatalf("Error creating Prometheus client: %v", err)
	}

	// Create a new v1 API instance
	v1api := v1.NewAPI(client)

	// Define the base query
	query := `sum(max_over_time(input_received_total{isUpdateConfig!="true",neosyncAccountId="8b34a902-0ec0-4612-b661-ba4d78b391c3"}[1d]))`
	// query = `max_over_time(input_received_total{neosyncAccountId="561b59fc-f582-4bdd-9e03-c6eb6ae27c5f"}[1d])` // stage
	// query = `max_over_time(input_received_total{neosyncJobId="d833b715-1118-4b7d-9489-3c0e1a5bd422"}[1d])`     // stage

	// Define the time range
	// end := time.Now()
	// start := end.AddDate(0, 0, -30) // 30 days ago

	start := time.Date(2024, 5, 1, 0, 0, 0, 0, time.UTC)
	end := time.Date(2024, 5, 31, 23, 59, 59, 0, time.UTC)

	// Collect daily results
	dailyResults, overallTotal, err := getDailyMaxValues(v1api, query, start, end)
	if err != nil {
		log.Fatalf("Error getting daily max values: %v", err)
	}

	// Print the daily maximum values
	fmt.Println("Daily Maximum Values:")
	for _, res := range dailyResults {
		fmt.Printf("Date: %s, Max Value: %.0f\n", res.Date.Format(usageDateFormat), res.Count)
	}

	// Print the overall total value
	fmt.Printf("Overall Total Value: %.0f\n", overallTotal)
}

func getDailyMaxValues(v1api v1.API, query string, start, end time.Time) ([]*DayResult, float64, error) {
	var dailyResults []*DayResult
	dailyTotals := make(map[string]float64)
	var overallTotal float64

	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	// Iterate through each day in the range
	i := 1
	errgrp, errctx := errgroup.WithContext(ctx)
	errgrp.SetLimit(10)
	mu := sync.Mutex{}

	fmt.Println(query, start, end)

	for d := start; !d.After(end); d = d.AddDate(0, 0, 1) {
		d := d
		errgrp.Go(func() error {
			dayStart := time.Date(d.Year(), d.Month(), d.Day(), 0, 0, 0, 0, time.UTC)
			dayEnd := dayStart.AddDate(0, 0, 1).Add(-time.Nanosecond)

			result, warnings, err := v1api.Query(errctx, query, dayEnd)
			if err != nil {
				return fmt.Errorf("error querying Prometheus for date %s: %w", d.Format(usageDateFormat), err)
			}
			if len(warnings) > 0 {
				log.Printf("Warnings for date %s: %v", d.Format(usageDateFormat), warnings)
			}

			// Process the results for this day
			vector, ok := result.(model.Vector)
			if !ok {
				return fmt.Errorf("error casting result to model.Vector for date %s", d.Format(usageDateFormat))
			}

			for _, sample := range vector {
				day := dayStart.Format(usageDateFormat)
				value := float64(sample.Value)
				mu.Lock()
				dailyTotals[day] += value
				mu.Unlock()
			}
			return nil
		})
		i++
	}
	err := errgrp.Wait()
	if err != nil {
		return nil, -1, err
	}

	var dates []string
	for day := range dailyTotals {
		dates = append(dates, day)
	}
	sort.Strings(dates)

	for _, day := range dates {
		date, err := time.Parse(usageDateFormat, day)
		if err != nil {
			return nil, 0, err
		}
		dailyResults = append(dailyResults, &DayResult{
			Date:  date,
			Count: dailyTotals[day],
		})
		overallTotal += dailyTotals[day]
	}

	return dailyResults, overallTotal, nil
}