nvllama-watchdog/exporters/influxdb.go

package exporters

import (
	"bytes"
	"context"
	"fmt"
	"net/http"
	"sort"
	"strings"
	"time"

	"git.govekar.net/LovroG05/nvllama-watchdog/core"
)

type InfluxDBExporter struct {
	url    string // full write URL, e.g. "http://localhost:8086/write?db=telegraf"
	token  string // for v2/v3
	user   string // for v1 basic auth
	pass   string // for v1 basic auth
	client *http.Client
}

// NewInfluxDBExporter creates an InfluxDB exporter.
//
// Version-specific URL examples:
//
//	v1: http://localhost:8086/write?db=nvllama&rp=autogen
//	v2: http://localhost:8086/api/v2/write?org=myorg&bucket=nvllama
//	v3: http://localhost:8086/api/v2/write?bucket=nvllama
func NewInfluxDBExporter(url, token, user, pass string) *InfluxDBExporter {
	return &InfluxDBExporter{
		url:    url,
		token:  token,
		user:   user,
		pass:   pass,
		client: &http.Client{Timeout: 10 * time.Second},
	}
}

func (e *InfluxDBExporter) Name() string { return "influxdb" }

func (e *InfluxDBExporter) Write(ctx context.Context, metrics []core.Metric) error {
	body := EncodeLineProtocol(metrics)
	if body == "" {
		return nil
	}

	req, err := http.NewRequestWithContext(ctx, http.MethodPost, e.url, bytes.NewReader([]byte(body)))
	if err != nil {
		return fmt.Errorf("influxdb: %w", err)
	}
	req.Header.Set("Content-Type", "text/plain; charset=utf-8")

	// Auth: token (v2/v3) or basic (v1)
	if e.token != "" {
		req.Header.Set("Authorization", "Token "+e.token)
	} else if e.user != "" {
		req.SetBasicAuth(e.user, e.pass)
	}

	resp, err := e.client.Do(req)
	if err != nil {
		return fmt.Errorf("influxdb: %w", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode >= 300 {
		buf := make([]byte, 256)
		n, _ := resp.Body.Read(buf)
		return fmt.Errorf("influxdb: HTTP %d: %s", resp.StatusCode, string(buf[:n]))
	}

	return nil
}

// ─── Line Protocol Encoder ──────────────────────────────────────────────────

// EncodeLineProtocol converts []core.Metric to InfluxDB line protocol.
//
// Each metric becomes one line:
//
//	measurement,tag1=val1,tag2=val2 value=<float> <timestamp_ns>
//
// The metric Name is the measurement, Labels are tags, and Value is
// written as a single "value" field.
func EncodeLineProtocol(metrics []core.Metric) string {
	if len(metrics) == 0 {
		return ""
	}

	// Group by (measurement + tags) so we can batch fields if needed.
	// For our use case each metric is one value, so one line per metric.
	var b strings.Builder
	b.Grow(len(metrics) * 128)

	for i, m := range metrics {
		if i > 0 {
			b.WriteByte('\n')
		}

		// Measurement (escape commas and spaces)
		b.WriteString(escapeMeasurement(m.Name))

		// Tags (sorted for determinism)
		if len(m.Labels) > 0 {
			keys := make([]string, 0, len(m.Labels))
			for k := range m.Labels {
				keys = append(keys, k)
			}
			sort.Strings(keys)

			for _, k := range keys {
				b.WriteByte(',')
				b.WriteString(escapeTagKey(k))
				b.WriteByte('=')
				b.WriteString(escapeTagValue(m.Labels[k]))
			}
		}

		// Fields
		b.WriteString(" value=")
		b.WriteString(formatFloat(m.Value))

		// Timestamp (nanoseconds since epoch)
		if !m.Timestamp.IsZero() {
			b.WriteByte(' ')
			fmt.Fprintf(&b, "%d", m.Timestamp.UnixNano())
		}
	}

	return b.String()
}

func escapeMeasurement(s string) string {
	// Escape commas and spaces in measurement names
	if strings.ContainsAny(s, ", ") {
		var b strings.Builder
		for _, c := range s {
			if c == ',' || c == ' ' {
				b.WriteByte('\\')
			}
			b.WriteRune(c)
		}
		return b.String()
	}
	return s
}

func escapeTagKey(s string) string {
	return escapeTagValue(s)
}

func escapeTagValue(s string) string {
	if strings.ContainsAny(s, ", =") {
		var b strings.Builder
		for _, c := range s {
			if c == ',' || c == '=' || c == ' ' {
				b.WriteByte('\\')
			}
			b.WriteRune(c)
		}
		return b.String()
	}
	return s
}

func formatFloat(f float64) string {
	// Write as integer if it's a whole number (InfluxDB treats unquoted
	// integers as i64, which is more storage-efficient)
	if f == float64(int64(f)) && f < 9.223372036854775807e18 && f > -9.223372036854775808e18 {
		return fmt.Sprintf("%di", int64(f))
	}
	return fmt.Sprintf("%g", f)
}