badger/main.go

package main

import (
	"context"
	"fmt"
	"image/color"
	"log"
	"net/http"
	"os"
	"path/filepath"
	"strconv"
	"time"

	"github.com/fogleman/gg"
	"github.com/gin-gonic/gin"
	"github.com/prometheus/client_golang/api"
	v1 "github.com/prometheus/client_golang/api/prometheus/v1"
	"github.com/prometheus/common/model"
)

const (
	// backgroundImageFilename = "images/bg.png"
	radius      = 0
	fontSize    = 12
	textMarginX = 5.0
	textMarginY = -2.0
)

var (
	queriesMap = map[string]queryParam{
		"mel":       {color.RGBA{49, 197, 83, 255}, "air_index_iaqi_pm25_v{city_name=\"Melbourne CBD\"}", "city_name"},
		"dockerhub": {color.RGBA{60, 174, 163, 255}, "docker_hub_pull_count{name=\"aria2-ariang\"}", "name"},
		"github":    {color.RGBA{32, 99, 155, 255}, "github_repository_stars{name=\"aria2-ariang-docker\"}", "name"},
	}
)

type Metric struct {
	Label string
	Value float64
	Time  time.Time
}

type queryParam struct {
	Color      color.Color
	Query      string
	ValueField string
}

func main() {
	client, err := api.NewClient(api.Config{
		Address: "https://prometheus-api.home.toozhao.com",
	})
	if err != nil {
		fmt.Printf("Error creating client: %v\n", err)
		os.Exit(1)
	}

	r := gin.Default()

	r.GET("/:name", func(c *gin.Context) {
		if err := generateBadge(client, c); err != nil {
			fmt.Println(err.Error())
		}
	})
	log.Fatal(r.Run())
}

func generateBadge(client api.Client, c *gin.Context) error {
	queryName := c.Param("name")
	query, ok := queriesMap[queryName]
	if !ok {
		query = queriesMap["mel"]
	}
	metric := queryMetric(client, query)

	dc := gg.NewContext(180, 22)
	fontPath := filepath.Join("fonts", "SourceCodePro-Regular.ttf")
	if err := dc.LoadFontFace(fontPath, fontSize); err != nil {
		return fmt.Errorf("load font %w", err)
	}

	// black layer
	dc.SetColor(color.RGBA{90, 90, 90, 255})
	dc.DrawRoundedRectangle(0, 0, float64(dc.Width()), float64(dc.Height()), radius)
	dc.Fill()

	labelWidth, textHeight := dc.MeasureString(metric.Label)

	// custom color layer
	dc.SetColor(query.Color)

	labelContainerWidth := labelWidth + textMarginX*2

	dc.DrawRectangle(0, 0, labelContainerWidth, float64(dc.Height()))
	dc.Fill()

	dc.SetColor(color.RGBA{255, 255, 255, 255})

	// label
	y := float64(dc.Height()) - textHeight - textMarginY
	dc.DrawString(metric.Label, textMarginX, y)

	// value
	valueText := strconv.Itoa(int(metric.Value))
	valueWidth, valueHeight := dc.MeasureString(valueText)
	dc.DrawString(valueText, float64(dc.Width())-valueWidth-textMarginX, float64(dc.Height())-valueHeight-textMarginY)
	w := c.Writer
	c.Status(http.StatusOK)
	w.Header().Set("Content-Type", "image/png")
	return dc.EncodePNG(w)
}

func queryMetric(client api.Client, query queryParam) *Metric {
	v1api := v1.NewAPI(client)
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	result, warnings, err := v1api.Query(ctx, query.Query, time.Now())
	if err != nil {
		fmt.Printf("Error querying Prometheus: %v\n", err)
		os.Exit(1)
	}
	if len(warnings) > 0 {
		fmt.Printf("Warnings: %v\n", warnings)
	}
	var metric Metric
	for _, v := range result.(model.Vector) {
		metric = Metric{
			Label: string(v.Metric[model.LabelName(query.ValueField)]),
			Value: float64(v.Value),
			Time:  v.Timestamp.Time(),
		}
	}
	return &metric
}