From 16330e41134b8351e009b99c4e8dc1a97ec968ad Mon Sep 17 00:00:00 2001
From: Carl Bergquist <carl.bergquist@gmail.com>
Date: Wed, 13 Nov 2024 14:36:12 +0100
Subject: [PATCH] Loki: Adds a `suggestions` resource call that inject scopes
 for finding label/labelvalues (#96025)

Signed-off-by: bergquist <carl.bergquist@gmail.com>
Co-authored-by: Sven Grossmann <sven.grossmann@grafana.com>
---
 pkg/tsdb/loki/loki.go   | 30 +++++++++++++----
 pkg/tsdb/loki/scopes.go | 71 +++++++++++++++++++++++++++++++++++++++++
 2 files changed, 95 insertions(+), 6 deletions(-)

diff --git a/pkg/tsdb/loki/loki.go b/pkg/tsdb/loki/loki.go
index a290c73ce1b..72ca44fe673 100644
--- a/pkg/tsdb/loki/loki.go
+++ b/pkg/tsdb/loki/loki.go
@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"net/http"
 	"regexp"
+	"strings"
 	"sync"
 	"time"
 
@@ -131,13 +132,30 @@ func callResource(ctx context.Context, req *backend.CallResourceRequest, sender
 
 	api := newLokiAPI(dsInfo.HTTPClient, dsInfo.URL, plog, tracer, false)
 
-	rawLokiResponse, err := api.RawQuery(ctx, lokiURL)
-	if err != nil {
-		span.RecordError(err)
-		span.SetStatus(codes.Error, err.Error())
-		plog.Error("Failed resource call from loki", "err", err, "url", lokiURL)
-		return err
+	var rawLokiResponse RawLokiResponse
+	var err error
+
+	// suggestions is a resource endpoint that will return label and label value suggestions based
+	// on queries and the existing scope. By moving this to the backend we can use the logql parser to
+	// rewrite queries safely.
+	if req.Method == http.MethodPost && strings.EqualFold(req.Path, "suggestions") {
+		rawLokiResponse, err = GetSuggestions(ctx, api, req)
+		if err != nil {
+			span.RecordError(err)
+			span.SetStatus(codes.Error, err.Error())
+			plog.FromContext(ctx).Error("Failed to get suggestions from loki", "err", err)
+			return err
+		}
+	} else {
+		rawLokiResponse, err = api.RawQuery(ctx, lokiURL)
+		if err != nil {
+			span.RecordError(err)
+			span.SetStatus(codes.Error, err.Error())
+			plog.Error("Failed resource call from loki", "err", err, "url", lokiURL)
+			return err
+		}
 	}
+
 	respHeaders := map[string][]string{
 		"content-type": {"application/json"},
 	}
diff --git a/pkg/tsdb/loki/scopes.go b/pkg/tsdb/loki/scopes.go
index c946cbb14fc..c9f502a780b 100644
--- a/pkg/tsdb/loki/scopes.go
+++ b/pkg/tsdb/loki/scopes.go
@@ -1,12 +1,83 @@
 package loki
 
 import (
+	"context"
+	"encoding/json"
 	"fmt"
+	"net/url"
+	"time"
 
+	"github.com/grafana/grafana-plugin-sdk-go/backend"
 	"github.com/grafana/grafana/pkg/promlib/models"
+	"github.com/grafana/grafana/pkg/tsdb/loki/kinds/dataquery"
 	"github.com/grafana/loki/v3/pkg/logql/syntax"
+	"github.com/prometheus/prometheus/promql/parser"
 )
 
+// SuggestionRequest is the request body for the GetSuggestions resource.
+type SuggestionRequest struct {
+	// LabelName, if provided, will result in label values being returned for the given label name.
+	LabelName string `json:"labelName"`
+
+	Query string `json:"query"`
+
+	Scopes       []models.ScopeFilter `json:"scopes"`
+	AdhocFilters []models.ScopeFilter `json:"adhocFilters"`
+
+	// Start and End are proxied directly to the prometheus endpoint (which is rfc3339 | unix_timestamp)
+	Start string `json:"start"`
+	End   string `json:"end"`
+}
+
+// GetSuggestions returns label names or label values for the given queries and scopes.
+func GetSuggestions(ctx context.Context, lokiAPI *LokiAPI, req *backend.CallResourceRequest) (RawLokiResponse, error) {
+	sugReq := SuggestionRequest{}
+	err := json.Unmarshal(req.Body, &sugReq)
+	if err != nil {
+		return RawLokiResponse{}, fmt.Errorf("error unmarshalling suggestion request: %v", err)
+	}
+
+	values := url.Values{}
+	if sugReq.Query != "" {
+		// the query is used to find label/labelvalues the duration and interval does not matter.
+		// If the user want to filter values based on time it should used the `start` and `end` fields
+		interpolatedQuery := interpolateVariables(sugReq.Query, time.Minute, time.Minute, dataquery.LokiQueryTypeRange, time.Minute)
+
+		if len(sugReq.Scopes) > 0 {
+			rewrittenQuery, err := ApplyScopes(interpolatedQuery, sugReq.Scopes)
+			if err == nil {
+				values.Add("query", rewrittenQuery)
+			} else {
+				values.Add("query", interpolatedQuery)
+			}
+		}
+	} else if len(sugReq.Scopes) > 0 {
+		matchers, err := models.FiltersToMatchers(sugReq.Scopes, sugReq.AdhocFilters)
+		if err != nil {
+			return RawLokiResponse{}, fmt.Errorf("error converting filters to matchers: %v", err)
+		}
+		vs := parser.VectorSelector{LabelMatchers: matchers}
+		values.Add("query", vs.String())
+	}
+
+	if sugReq.Start != "" {
+		values.Add("start", sugReq.Start)
+	}
+
+	if sugReq.End != "" {
+		values.Add("end", sugReq.End)
+	}
+
+	var path string
+	if sugReq.LabelName != "" {
+		path = "/loki/api/v1/label/" + url.QueryEscape(sugReq.LabelName) + "/values?" + values.Encode()
+	} else {
+		path = "/loki/api/v1/labels?" + values.Encode()
+	}
+
+	return lokiAPI.RawQuery(ctx, path)
+}
+
 // ApplyScopes applies the given scope filters to the given raw expression.
 func ApplyScopes(rawExpr string, scopeFilters []models.ScopeFilter) (string, error) {
 	if len(scopeFilters) == 0 {