labmcp/internal/monitoring/loki.go

package monitoring

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"strings"
	"time"
)

// LokiClient is an HTTP client for the Loki API.
type LokiClient struct {
	baseURL    string
	httpClient *http.Client
}

// NewLokiClient creates a new Loki API client.
func NewLokiClient(baseURL string) *LokiClient {
	return &LokiClient{
		baseURL: strings.TrimRight(baseURL, "/"),
		httpClient: &http.Client{
			Timeout: 30 * time.Second,
		},
	}
}

// QueryRange executes a LogQL range query against Loki.
func (c *LokiClient) QueryRange(ctx context.Context, logql string, start, end time.Time, limit int, direction string) (*LokiQueryData, error) {
	params := url.Values{}
	params.Set("query", logql)
	params.Set("start", fmt.Sprintf("%d", start.UnixNano()))
	params.Set("end", fmt.Sprintf("%d", end.UnixNano()))
	if limit > 0 {
		params.Set("limit", fmt.Sprintf("%d", limit))
	}
	if direction != "" {
		params.Set("direction", direction)
	}

	body, err := c.get(ctx, "/loki/api/v1/query_range", params)
	if err != nil {
		return nil, fmt.Errorf("query range failed: %w", err)
	}

	var data LokiQueryData
	if err := json.Unmarshal(body, &data); err != nil {
		return nil, fmt.Errorf("failed to parse query data: %w", err)
	}
	return &data, nil
}

// Labels returns all available label names from Loki.
func (c *LokiClient) Labels(ctx context.Context) ([]string, error) {
	body, err := c.get(ctx, "/loki/api/v1/labels", nil)
	if err != nil {
		return nil, fmt.Errorf("labels failed: %w", err)
	}

	var labels []string
	if err := json.Unmarshal(body, &labels); err != nil {
		return nil, fmt.Errorf("failed to parse labels: %w", err)
	}
	return labels, nil
}

// LabelValues returns all values for a given label name from Loki.
func (c *LokiClient) LabelValues(ctx context.Context, label string) ([]string, error) {
	path := fmt.Sprintf("/loki/api/v1/label/%s/values", url.PathEscape(label))
	body, err := c.get(ctx, path, nil)
	if err != nil {
		return nil, fmt.Errorf("label values failed: %w", err)
	}

	var values []string
	if err := json.Unmarshal(body, &values); err != nil {
		return nil, fmt.Errorf("failed to parse label values: %w", err)
	}
	return values, nil
}

// get performs a GET request and returns the "data" field from the Loki response envelope.
// Loki uses the same {"status":"success","data":...} format as Prometheus.
func (c *LokiClient) get(ctx context.Context, path string, params url.Values) (json.RawMessage, error) {
	u := c.baseURL + path
	if len(params) > 0 {
		u += "?" + params.Encode()
	}

	req, err := http.NewRequestWithContext(ctx, http.MethodGet, u, nil)
	if err != nil {
		return nil, fmt.Errorf("failed to create request: %w", err)
	}

	resp, err := c.httpClient.Do(req)
	if err != nil {
		return nil, fmt.Errorf("request failed: %w", err)
	}
	defer resp.Body.Close() //nolint:errcheck // cleanup on exit

	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, fmt.Errorf("failed to read response: %w", err)
	}

	if resp.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("unexpected status %d: %s", resp.StatusCode, string(body))
	}

	var promResp PromResponse
	if err := json.Unmarshal(body, &promResp); err != nil {
		return nil, fmt.Errorf("failed to parse response: %w", err)
	}

	if promResp.Status != "success" {
		return nil, fmt.Errorf("loki error (%s): %s", promResp.ErrorType, promResp.Error)
	}

	return promResp.Data, nil
}