nmon/internal/api/query.go

package api

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"net/http"
	"strconv"
	"strings"
	"time"

	"github.com/danielgtaylor/huma/v2"
)

var jsonObject = map[string]*huma.Response{
	"200": {Description: "The store's Prometheus-compatible JSON response, unchanged.", Content: map[string]*huma.MediaType{
		"application/json": {Schema: &huma.Schema{Type: "object"}},
	}},
}

func (s *server) registerQuery(api huma.API) {
	huma.Register(api, huma.Operation{
		OperationID: "query", Method: "GET", Path: "/api/v1/query", Tags: []string{"query"},
		Summary:     "Instant MetricsQL query against the hot store",
		Description: "Read-only proxy with a 30 s timeout. A result of more than 100,000 series is rejected with 422.",
		Responses:   jsonObject, Errors: []int{422, 503, 504},
	}, s.query)
	huma.Register(api, huma.Operation{
		OperationID: "query-range", Method: "GET", Path: "/api/v1/query_range", Tags: []string{"query"},
		Summary:     "Range MetricsQL query against the hot store",
		Description: "Same limits as /api/v1/query.",
		Responses:   jsonObject, Errors: []int{422, 503, 504},
	}, s.queryRange)
}

func raw(body []byte) *huma.StreamResponse {
	return &huma.StreamResponse{Body: func(ctx huma.Context) {
		ctx.SetHeader("Content-Type", "application/json")
		ctx.SetStatus(http.StatusOK)
		_, _ = ctx.BodyWriter().Write(body)
	}}
}

func (s *server) query(ctx context.Context, in *struct {
	Q    string `query:"q" required:"true" minLength:"1" doc:"MetricsQL expression." example:"ifHCInOctets"`
	Time string `query:"time" doc:"Evaluation time: RFC 3339 or Unix seconds. Default now."`
}) (*huma.StreamResponse, error) {
	var ts time.Time
	if in.Time != "" {
		var err error
		if ts, err = parseTime(in.Time); err != nil {
			return nil, huma.Error422UnprocessableEntity("time: " + err.Error())
		}
	}
	release, err := s.admit(ctx, true)
	if err != nil {
		return nil, err
	}
	defer release()
	qctx, cancel := context.WithTimeout(ctx, QueryTimeout)
	defer cancel()
	body, err := s.Querier.Query(qctx, "hot", in.Q, ts)
	return s.finish(qctx, body, err)
}

func (s *server) queryRange(ctx context.Context, in *struct {
	Q    string `query:"q" required:"true" minLength:"1" doc:"MetricsQL expression."`
	From string `query:"from" doc:"Start: RFC 3339 or Unix seconds. Default one hour before to."`
	To   string `query:"to" doc:"End: RFC 3339 or Unix seconds. Default now."`
	Step string `query:"step" required:"true" doc:"Resolution step: a duration such as 30s or a number of seconds."`
}) (*huma.StreamResponse, error) {
	to := s.Clock.Now()
	if in.To != "" {
		var err error
		if to, err = parseTime(in.To); err != nil {
			return nil, huma.Error422UnprocessableEntity("to: " + err.Error())
		}
	}
	from := to.Add(-time.Hour)
	if in.From != "" {
		var err error
		if from, err = parseTime(in.From); err != nil {
			return nil, huma.Error422UnprocessableEntity("from: " + err.Error())
		}
	}
	step, err := parseStep(in.Step)
	if err != nil {
		return nil, huma.Error422UnprocessableEntity("step: " + err.Error())
	}
	if !from.Before(to) {
		return nil, huma.Error422UnprocessableEntity("from must be before to")
	}
	release, err := s.admit(ctx, true)
	if err != nil {
		return nil, err
	}
	defer release()
	qctx, cancel := context.WithTimeout(ctx, QueryTimeout)
	defer cancel()
	body, err := s.Querier.QueryRange(qctx, "hot", in.Q, from, to, step)
	return s.finish(qctx, body, err)
}

// finish applies the series cap and maps store errors to problems.
func (s *server) finish(ctx context.Context, body []byte, err error) (*huma.StreamResponse, error) {
	if err != nil {
		return nil, s.vmErr(ctx, err)
	}
	if n := countSeries(body); n > MaxSeries {
		return nil, huma.Error422UnprocessableEntity(
			"query returned " + itoa(n) + " series; the limit is " + itoa(MaxSeries) + ", narrow the selector")
	}
	return raw(body), nil
}

func errStoreDown() error {
	return huma.Error503ServiceUnavailable("the metric store is not available")
}

func errQueryTimeout() error {
	return huma.Error504GatewayTimeout("query exceeded " + QueryTimeout.String())
}

func errStoreFailed() error { return huma.Error502BadGateway("metric store error") }

// isBadQuery reports a rejection by the store of the query itself.
func isBadQuery(err error) bool {
	return strings.Contains(err.Error(), "400 Bad Request") || strings.Contains(err.Error(), "422 Unprocessable")
}

func badQuery(err error) error { return huma.Error422UnprocessableEntity(err.Error()) }

// countSeries returns len(data.result) of a Prometheus-style response.
func countSeries(body []byte) int {
	var r struct {
		Data struct {
			Result []json.RawMessage `json:"result"`
		} `json:"data"`
	}
	if json.NewDecoder(bytes.NewReader(body)).Decode(&r) != nil {
		return 0
	}
	return len(r.Data.Result)
}

func parseTime(s string) (time.Time, error) {
	if t, err := time.Parse(time.RFC3339, s); err == nil {
		return t.UTC(), nil
	}
	f, err := strconv.ParseFloat(s, 64)
	if err != nil {
		return time.Time{}, errors.New("want RFC 3339 or Unix seconds")
	}
	return time.UnixMilli(int64(f * 1000)).UTC(), nil
}

func parseStep(s string) (time.Duration, error) {
	d, err := time.ParseDuration(s)
	if err != nil {
		f, ferr := strconv.ParseFloat(s, 64)
		if ferr != nil {
			return 0, errors.New("want a duration such as 30s or a number of seconds")
		}
		d = time.Duration(f * float64(time.Second))
	}
	if d <= 0 {
		return 0, errors.New("must be positive")
	}
	return d, nil
}

func (s *server) registerMetrics(api huma.API) {
	huma.Register(api, huma.Operation{
		OperationID: "metrics", Method: "GET", Path: "/metrics", Tags: []string{"internals"},
		Summary: "nmon's own metrics in Prometheus text format",
		Responses: map[string]*huma.Response{"200": {Description: "Prometheus text exposition.", Content: map[string]*huma.MediaType{
			"text/plain": {Schema: &huma.Schema{Type: "string"}},
		}}},
	}, func(context.Context, *struct{}) (*huma.StreamResponse, error) {
		return &huma.StreamResponse{Body: func(ctx huma.Context) {
			ctx.SetHeader("Content-Type", "text/plain; version=0.0.4; charset=utf-8")
			_ = s.Registry.WritePrometheus(ctx.BodyWriter())
		}}, nil
	})
}