query-orchestration/internal/serviceconfig/observability/prometheus/generator/main.go

//ignore go:build ignore

package main

import (
	"bytes"
	"fmt"
	"go/format"
	"log"
	"os"
	"strings"
	"text/template"

	"golang.org/x/text/cases"
	"golang.org/x/text/language"

	"queryorchestration/internal/serviceconfig/observability/prometheus"
)

const outputFilename = "./metric_stubs.gen.go"

// funcTemplate is the template for generating each metric recording function
const funcTemplate = `
// Record{{.FuncName}} records a {{.Name}} metric
{{.Comment}}
func (m *Metrics) Record{{.FuncName}}({{.Params}}) error {
    return m.RecordMetric("{{.Name}}", {{.DefaultValue}}, map[string]string{ {{.LabelMap}} })
}`

// templateData holds the data needed for the template
type templateData struct {
	FuncName     string
	Name         string
	Params       string
	DefaultValue string
	LabelMap     string
	Comment      string
}

// reservedKeywords contains Go reserved keywords that can't be used as parameter names
var reservedKeywords = map[string]bool{
	"type":   true,
	"range":  true,
	"map":    true,
	"func":   true,
	"return": true,
	"goto":   true,
	"break":  true,
}

// toCamelCase function
func toCamelCase(s string) string {
	parts := strings.Split(s, "_")
	casesTitle := cases.Title(language.English)
	for i := range parts {
		parts[i] = casesTitle.String(parts[i])
	}
	return strings.Join(parts, "")
}

// getSafeParamName returns a safe parameter name that avoids Go reserved keywords
func getSafeParamName(label string) string {
	paramName := strings.Replace(label, "-", "", -1)
	if reservedKeywords[paramName] {
		return paramName + "Value"
	}
	return paramName
}

func generateFunctionData(metric prometheus.MetricDefinition) templateData {
	// Generate function name
	funcName := toCamelCase(metric.Name)

	// Prepare parameters
	var params []string

	// Always add value parameter since RecordMetric always needs it
	params = append(params, "value float64")

	// Add label parameters
	for _, label := range metric.Labels {
		paramName := getSafeParamName(label)
		params = append(params, fmt.Sprintf("%s string", paramName))
	}

	// Generate label map
	var labelPairs []string
	for _, label := range metric.Labels {
		paramName := getSafeParamName(label)
		labelPairs = append(labelPairs, fmt.Sprintf(`"%s": %s`, label, paramName))
	}

	// Generate comment with parameter documentation
	var commentLines []string
	if metric.Help != "" {
		commentLines = append(commentLines, fmt.Sprintf("// %s", metric.Help))
	}
	commentLines = append(commentLines, "//")
	commentLines = append(commentLines, "// Parameters:")
	commentLines = append(commentLines, "//   value: The value to record")
	for _, label := range metric.Labels {
		paramName := getSafeParamName(label)
		commentLines = append(commentLines, fmt.Sprintf("//   %s: The %s label value", paramName, label))
	}

	return templateData{
		FuncName:     funcName,
		Name:         metric.Name,
		Params:       strings.Join(params, ", "),
		DefaultValue: "value", // Always use the value parameter
		LabelMap:     strings.Join(labelPairs, ", "),
		Comment:      strings.Join(commentLines, "\n"),
	}
}

func main() {
	// Get metric definitions from the package above
	metrics := prometheus.NewStandardMetricDefinitions()

	// Create the template
	tmpl, err := template.New("function").Parse(funcTemplate)
	if err != nil {
		log.Fatalf("Error creating template: %v", err)
	}

	// Generate the output file
	var buf bytes.Buffer

	// Write package declaration and imports
	buf.WriteString("// Code generated by metric-generator; DO NOT EDIT.\n\n")
	buf.WriteString("package prometheus\n\n")

	// Generate functions for each metric
	for _, metric := range metrics {
		data := generateFunctionData(metric)

		if err := tmpl.Execute(&buf, data); err != nil {
			log.Fatalf("Error executing template for %s: %v", metric.Name, err)
		}
		buf.WriteString("\n\n") // Add spacing between functions
	}

	// Format the generated code
	formattedCode, err := format.Source(buf.Bytes())
	if err != nil {
		log.Fatalf("Error formatting code: %v", err)
	}

	// Write to file
	if err := os.WriteFile(outputFilename, formattedCode, 0644); err != nil {
		log.Fatalf("Error writing output file: %v", err)
	}

	fmt.Println("Successfully generated " + outputFilename)
}