Add dedicated instant/range query handlers

Signed-off-by: SungJin1212 <tjdwls1201@gmail.com>

Author: SungJin1212

CHANGELOG.md

@@ -14,6 +14,7 @@
 * [FEATURE] Experimental Support Parquet format: Implement parquet converter service to convert a TSDB block into Parquet and Parquet Queryable. #6716 #6743
 * [FEATURE] Distributor/Ingester: Implemented experimental feature to use gRPC stream connection for push requests. This can be enabled by setting `-distributor.use-stream-push=true`. #6580
 * [FEATURE] Compactor: Add support for percentage based sharding for compactors. #6738
+* [FEATURE] Querier: Add dedicated instant/range query handlers to customize these APIs. #6763
 * [ENHANCEMENT] Query Frontend: Change to return 400 when the tenant resolving fail. #6715
 * [ENHANCEMENT] Querier: Support query parameters to metadata api (/api/v1/metadata) to allow user to limit metadata to return. Add a `-ingester.return-all-metadata` flag to make the metadata API run when the deployment. Please set this flag to `false` to use the metadata API with the limits later. #6681 #6744
 * [ENHANCEMENT] Ingester: Add a `cortex_ingester_active_native_histogram_series` metric to track # of active NH series. #6695

pkg/api/custom_api.go

@@ -0,0 +1,298 @@
+package api
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"net/http"
+	"path"
+	"time"
+
+	"github.com/go-kit/log"
+	"github.com/go-kit/log/level"
+	"github.com/gorilla/mux"
+	"github.com/grafana/regexp"
+	jsoniter "github.com/json-iterator/go"
+	"github.com/munnerz/goautoneg"
+	"github.com/prometheus/prometheus/promql"
+	"github.com/prometheus/prometheus/storage"
+	"github.com/prometheus/prometheus/util/annotations"
+	"github.com/prometheus/prometheus/util/httputil"
+	v1 "github.com/prometheus/prometheus/web/api/v1"
+)
+
+type CustomAPI struct {
+	queryable     storage.SampleAndChunkQueryable
+	queryEngine   promql.QueryEngine
+	now           func() time.Time
+	statsRenderer v1.StatsRenderer
+	logger        log.Logger
+	codecs        []v1.Codec
+	CORSOrigin    *regexp.Regexp
+}
+
+func NewCustomAPI(
+	qe promql.QueryEngine,
+	q storage.SampleAndChunkQueryable,
+	statsRenderer v1.StatsRenderer,
+	logger log.Logger,
+	codecs []v1.Codec,
+	CORSOrigin *regexp.Regexp,
+) *CustomAPI {
+	return &CustomAPI{
+		queryable:     q,
+		queryEngine:   qe,
+		now:           time.Now,
+		statsRenderer: statsRenderer,
+		logger:        logger,
+		codecs:        codecs,
+		CORSOrigin:    CORSOrigin,
+	}
+}
+
+// Register custom api path with handler
+// Usage: router.Path(path.Join(prefix, ${path})).Methods(${method}).Handler(${handler})
+func (c *CustomAPI) Register(router *mux.Router, prefix string) {
+	router.Path(path.Join(prefix, "/api/v1/query")).Methods("GET", "POST").Handler(c.wrap(c.instantHandler))
+	router.Path(path.Join(prefix, "/api/v1/query_range")).Methods("GET", "POST").Handler(c.wrap(c.rangeQueryHandler))
+}
+
+// Custom handler for Query range API
+func (c *CustomAPI) rangeQueryHandler(r *http.Request) (result apiFuncResult) {
+	start, err := parseTime(r.FormValue("start"))
+	if err != nil {
+		return invalidParamError(err, "start")
+	}
+	end, err := parseTime(r.FormValue("end"))
+	if err != nil {
+		return invalidParamError(err, "end")
+	}
+	if end.Before(start) {
+		return invalidParamError(errors.New("end timestamp must not be before start time"), "end")
+	}
+
+	step, err := parseDuration(r.FormValue("step"))
+	if err != nil {
+		return invalidParamError(err, "step")
+	}
+
+	if step <= 0 {
+		return invalidParamError(errors.New("zero or negative query resolution step widths are not accepted. Try a positive integer"), "step")
+	}
+
+	// For safety, limit the number of returned points per timeseries.
+	// This is sufficient for 60s resolution for a week or 1h resolution for a year.
+	if end.Sub(start)/step > 11000 {
+		err := errors.New("exceeded maximum resolution of 11,000 points per timeseries. Try decreasing the query resolution (?step=XX)")
+		return apiFuncResult{nil, &apiError{errorBadData, err}, nil, nil}
+	}
+
+	ctx := r.Context()
+	if to := r.FormValue("timeout"); to != "" {
+		var cancel context.CancelFunc
+		timeout, err := parseDuration(to)
+		if err != nil {
+			return invalidParamError(err, "timeout")
+		}
+
+		ctx, cancel = context.WithTimeout(ctx, timeout)
+		defer cancel()
+	}
+
+	opts, err := extractQueryOpts(r)
+	if err != nil {
+		return apiFuncResult{nil, &apiError{errorBadData, err}, nil, nil}
+	}
+	qry, err := c.queryEngine.NewRangeQuery(ctx, c.queryable, opts, r.FormValue("query"), start, end, step)
+	if err != nil {
+		return invalidParamError(err, "query")
+	}
+	// From now on, we must only return with a finalizer in the result (to
+	// be called by the caller) or call qry.Close ourselves (which is
+	// required in the case of a panic).
+	defer func() {
+		if result.finalizer == nil {
+			qry.Close()
+		}
+	}()
+
+	ctx = httputil.ContextFromRequest(ctx, r)
+
+	res := qry.Exec(ctx)
+	if res.Err != nil {
+		return apiFuncResult{nil, returnAPIError(res.Err), res.Warnings, qry.Close}
+	}
+
+	warnings := res.Warnings
+	qs := c.statsRenderer(ctx, qry.Stats(), r.FormValue("stats"))
+
+	return apiFuncResult{&v1.QueryData{
+		ResultType: res.Value.Type(),
+		Result:     res.Value,
+		Stats:      qs,
+	}, nil, warnings, qry.Close}
+}
+
+// Custom handler for Query API
+func (c *CustomAPI) instantHandler(r *http.Request) (result apiFuncResult) {
+	ts, err := parseTimeParam(r, "time", c.now())
+	if err != nil {
+		return invalidParamError(err, "time")
+	}
+
+	ctx := r.Context()
+	if to := r.FormValue("timeout"); to != "" {
+		var cancel context.CancelFunc
+		timeout, err := parseDuration(to)
+		if err != nil {
+			return invalidParamError(err, "timeout")
+		}
+
+		ctx, cancel = context.WithDeadline(ctx, c.now().Add(timeout))
+		defer cancel()
+	}
+
+	opts, err := extractQueryOpts(r)
+	if err != nil {
+		return apiFuncResult{nil, &apiError{errorBadData, err}, nil, nil}
+	}
+	qry, err := c.queryEngine.NewInstantQuery(ctx, c.queryable, opts, r.FormValue("query"), ts)
+	if err != nil {
+		return invalidParamError(err, "query")
+	}
+
+	// From now on, we must only return with a finalizer in the result (to
+	// be called by the caller) or call qry.Close ourselves (which is
+	// required in the case of a panic).
+	defer func() {
+		if result.finalizer == nil {
+			qry.Close()
+		}
+	}()
+
+	ctx = httputil.ContextFromRequest(ctx, r)
+
+	res := qry.Exec(ctx)
+	if res.Err != nil {
+		return apiFuncResult{nil, returnAPIError(res.Err), res.Warnings, qry.Close}
+	}
+
+	warnings := res.Warnings
+	qs := c.statsRenderer(ctx, qry.Stats(), r.FormValue("stats"))
+
+	return apiFuncResult{&v1.QueryData{
+		ResultType: res.Value.Type(),
+		Result:     res.Value,
+		Stats:      qs,
+	}, nil, warnings, qry.Close}
+}
+
+func (c *CustomAPI) wrap(f apiFunc) http.HandlerFunc {
+	return func(w http.ResponseWriter, r *http.Request) {
+		httputil.SetCORS(w, c.CORSOrigin, r)
+
+		result := f(r)
+		if result.finalizer != nil {
+			defer result.finalizer()
+		}
+
+		if result.err != nil {
+			c.respondError(w, result.err, result.data)
+			return
+		}
+
+		if result.data != nil {
+			c.respond(w, r, result.data, result.warnings, r.FormValue("query"))
+			return
+		}
+		w.WriteHeader(http.StatusNoContent)
+	}
+}
+
+func (c *CustomAPI) respondError(w http.ResponseWriter, apiErr *apiError, data interface{}) {
+	json := jsoniter.ConfigCompatibleWithStandardLibrary
+	b, err := json.Marshal(&response{
+		Status:    statusError,
+		ErrorType: apiErr.typ,
+		Error:     apiErr.err.Error(),
+		Data:      data,
+	})
+	if err != nil {
+		level.Error(c.logger).Log("error marshaling json response", "err", err)
+		http.Error(w, err.Error(), http.StatusInternalServerError)
+		return
+	}
+
+	var code int
+	switch apiErr.typ {
+	case errorBadData:
+		code = http.StatusBadRequest
+	case errorExec:
+		code = http.StatusUnprocessableEntity
+	case errorCanceled:
+		code = statusClientClosedConnection
+	case errorTimeout:
+		code = http.StatusServiceUnavailable
+	case errorInternal:
+		code = http.StatusInternalServerError
+	case errorNotFound:
+		code = http.StatusNotFound
+	case errorNotAcceptable:
+		code = http.StatusNotAcceptable
+	default:
+		code = http.StatusInternalServerError
+	}
+
+	w.Header().Set("Content-Type", "application/json")
+	w.WriteHeader(code)
+	if n, err := w.Write(b); err != nil {
+		level.Error(c.logger).Log("error writing response", "bytesWritten", n, "err", err)
+	}
+}
+
+func (c *CustomAPI) respond(w http.ResponseWriter, req *http.Request, data interface{}, warnings annotations.Annotations, query string) {
+	warn, info := warnings.AsStrings(query, 10, 10)
+
+	resp := &v1.Response{
+		Status:   statusSuccess,
+		Data:     data,
+		Warnings: warn,
+		Infos:    info,
+	}
+
+	codec, err := c.negotiateCodec(req, resp)
+	if err != nil {
+		c.respondError(w, &apiError{errorNotAcceptable, err}, nil)
+		return
+	}
+
+	b, err := codec.Encode(resp)
+	if err != nil {
+		level.Error(c.logger).Log("error marshaling response", "url", req.URL, "err", err)
+		http.Error(w, err.Error(), http.StatusInternalServerError)
+		return
+	}
+
+	w.Header().Set("Content-Type", codec.ContentType().String())
+	w.WriteHeader(http.StatusOK)
+	if n, err := w.Write(b); err != nil {
+		level.Error(c.logger).Log("error writing response", "url", req.URL, "bytesWritten", n, "err", err)
+	}
+}
+
+func (c *CustomAPI) negotiateCodec(req *http.Request, resp *v1.Response) (v1.Codec, error) {
+	for _, clause := range goautoneg.ParseAccept(req.Header.Get("Accept")) {
+		for _, codec := range c.codecs {
+			if codec.ContentType().Satisfies(clause) && codec.CanEncode(resp) {
+				return codec, nil
+			}
+		}
+	}
+
+	defaultCodec := c.codecs[0]
+	if !defaultCodec.CanEncode(resp) {
+		return nil, fmt.Errorf("cannot encode response as %s", defaultCodec.ContentType())
+	}
+
+	return defaultCodec, nil
+}