public/grafana
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Loki: Set Query limits context header (#114112)

Browse Source 
* feat(lokiQueryLimitsContext): send full query range & expr on initial split requests
This commit is contained in:
Galen Kistler
2025-11-21 07:19:55 -06:00
committed by GitHub
parent 669382c212
commit c0f6d90971
26 changed files with 1088 additions and 1335 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pkg/tsdb/loki/api.go
+11
View File
@@ -96,6 +96,8 @@ func makeDataRequest(ctx context.Context, lokiDsUrl string, query lokiQuery) (*h
return nil, backend.DownstreamError(fmt.Errorf("failed to create request: %w", err))
}
addQueryLimitsHeader(query, req)
if query.SupportingQueryType != SupportingQueryNone {
value := getSupportingQueryHeaderValue(query.SupportingQueryType)
if value != "" {
@@ -108,6 +110,15 @@ func makeDataRequest(ctx context.Context, lokiDsUrl string, query lokiQuery) (*h
return req, nil
}
func addQueryLimitsHeader(query lokiQuery, req *http.Request) {
if len(query.LimitsContext.Expr) > 0 {
queryLimitStr, err := json.Marshal(query.LimitsContext)
if err == nil {
req.Header.Set("X-Loki-Query-Limits-Context", string(queryLimitStr))
}
}
}
type lokiResponseError struct {
Message string `json:"message"`
TraceID string `json:"traceID,omitempty"`
pkg/tsdb/loki/api_test.go
+52
View File
@@ -2,10 +2,12 @@ package loki
import (
"context"
"encoding/json"
"fmt"
"net/http"
"strings"
"testing"
"time"
"github.com/grafana/grafana-plugin-sdk-go/backend"
"github.com/grafana/grafana/pkg/tsdb/loki/kinds/dataquery"
@@ -47,6 +49,56 @@ func TestApiLogVolume(t *testing.T) {
require.True(t, called)
})
t.Run("X-Loki-Query-Limits-Context header should be set when LimitsContext is provided", func(t *testing.T) {
called := false
from := time.Now().Truncate(time.Millisecond).Add(-1 * time.Hour)
to := time.Now().Truncate(time.Millisecond)
limitsContext := LimitsContext{
Expr: "{cluster=\"us-central1\"}",
From: from,
To: to,
}
limitsContextJson, _ := json.Marshal(limitsContext)
api := makeMockedAPI(200, "application/json", response, func(req *http.Request) {
called = true
require.Equal(t, string(limitsContextJson), req.Header.Get("X-Loki-Query-Limits-Context"))
})
_, err := api.DataQuery(context.Background(), lokiQuery{Expr: "", SupportingQueryType: SupportingQueryLogsSample, QueryType: QueryTypeRange, LimitsContext: limitsContext}, ResponseOpts{})
require.NoError(t, err)
require.True(t, called)
})
t.Run("X-Loki-Query-Limits-Context header should not get set when LimitsContext is missing expr", func(t *testing.T) {
called := false
from := time.Now().Truncate(time.Millisecond).Add(-1 * time.Hour)
to := time.Now().Truncate(time.Millisecond)
limitsContext := LimitsContext{
Expr: "",
From: from,
To: to,
}
api := makeMockedAPI(200, "application/json", response, func(req *http.Request) {
called = true
require.Equal(t, "", req.Header.Get("X-Loki-Query-Limits-Context"))
})
_, err := api.DataQuery(context.Background(), lokiQuery{Expr: "", SupportingQueryType: SupportingQueryLogsSample, QueryType: QueryTypeRange, LimitsContext: limitsContext}, ResponseOpts{})
require.NoError(t, err)
require.True(t, called)
})
t.Run("X-Loki-Query-Limits-Context header should not get set when LimitsContext is not provided", func(t *testing.T) {
called := false
api := makeMockedAPI(200, "application/json", response, func(req *http.Request) {
called = true
require.Equal(t, "", req.Header.Get("X-Loki-Query-Limits-Context"))
})
_, err := api.DataQuery(context.Background(), lokiQuery{Expr: "", SupportingQueryType: SupportingQueryLogsSample, QueryType: QueryTypeRange}, ResponseOpts{})
require.NoError(t, err)
require.True(t, called)
})
t.Run("data sample queries should set data sample http header", func(t *testing.T) {
called := false
api := makeMockedAPI(200, "application/json", response, func(req *http.Request) {
pkg/tsdb/loki/kinds/dataquery/types_dataquery_gen.go
Generated
+13
View File
@@ -18,6 +18,17 @@ const (
QueryEditorModeBuilder QueryEditorMode = "builder"
)
type LimitsContext struct {
Expr string `json:"expr"`
From int64 `json:"from"`
To int64 `json:"to"`
}
// NewLimitsContext creates a new LimitsContext object.
func NewLimitsContext() *LimitsContext {
return &LimitsContext{}
}
type LokiQueryType string
const (
@@ -59,6 +70,8 @@ type LokiDataQuery struct {
Instant *bool `json:"instant,omitempty"`
// Used to set step value for range queries.
Step *string `json:"step,omitempty"`
// The full query plan for split/shard queries. Encoded and sent to Loki via `X-Loki-Query-Limits-Context` header. Requires "lokiQueryLimitsContext" feature flag
LimitsContext *LimitsContext `json:"limitsContext,omitempty"`
// A unique identifier for the query within the list of targets.
// In server side expressions, the refId is used as a variable name to identify results.
// By default, the UI will assign A->Z; however setting meaningful names may be useful.
pkg/tsdb/loki/parse_query.go
+15
View File
@@ -156,6 +156,8 @@ func parseQuery(queryContext *backend.QueryDataRequest, logqlScopesEnabled bool)
expr := interpolateVariables(model.Expr, interval, timeRange, queryType, step)
limitsConfig := generateLimitsConfig(model, interval, timeRange, queryType, step)
direction, err := parseDirection(model.Direction)
if err != nil {
return nil, err
@@ -192,8 +194,21 @@ func parseQuery(queryContext *backend.QueryDataRequest, logqlScopesEnabled bool)
RefID: query.RefID,
SupportingQueryType: supportingQueryType,
Scopes: model.Scopes,
LimitsContext: limitsConfig,
})
}
return qs, nil
}
func generateLimitsConfig(model *QueryJSONModel, interval time.Duration, timeRange time.Duration, queryType QueryType, step time.Duration) LimitsContext {
var limitsConfig LimitsContext
// Only supply limits context config if we have expression, and from and to
if model.LimitsContext != nil && model.LimitsContext.Expr != "" && model.LimitsContext.From > 0 && model.LimitsContext.To > 0 {
// If a limits expression was provided, interpolate it and parse the time range
limitsConfig.Expr = interpolateVariables(model.LimitsContext.Expr, interval, timeRange, queryType, step)
limitsConfig.From = time.UnixMilli(model.LimitsContext.From)
limitsConfig.To = time.UnixMilli(model.LimitsContext.To)
}
return limitsConfig
}
pkg/tsdb/loki/parse_query_test.go
+69
View File
@@ -1,6 +1,7 @@
package loki
import (
"strconv"
"testing"
"time"
@@ -145,6 +146,74 @@ func TestParseQuery(t *testing.T) {
require.Equal(t, `{namespace="logish"} |= "problems"`, models[0].Expr)
})
t.Run("parsing query model with invalid query limits context expr", func(t *testing.T) {
from := time.Now().Add(-3000 * time.Second)
fullFrom := time.Now().Add(-1 * time.Hour)
to := time.Now()
queryContext := &backend.QueryDataRequest{
Queries: []backend.DataQuery{
{
JSON: []byte(`
{
"expr": "count_over_time({service_name=\"apache\", __stream_shard__=\"2\"}[$__auto])",
"format": "time_series",
"refId": "A",
"limitsContext": {"expr": "", "from": ` + strconv.FormatInt(fullFrom.UnixMilli(), 10) + `, "to": ` + strconv.FormatInt(to.UnixMilli(), 10) + `}
}`,
),
TimeRange: backend.TimeRange{
From: from,
To: to,
},
Interval: time.Second * 15,
MaxDataPoints: 200,
},
},
}
models, err := parseQuery(queryContext, true)
require.NoError(t, err)
require.Equal(t, `count_over_time({service_name="apache", __stream_shard__="2"}[15s])`, models[0].Expr)
// If the limits context expression is missing, we don't set any limits context
require.Equal(t, ``, models[0].LimitsContext.Expr)
require.Equal(t, time.Time{}, models[0].LimitsContext.To)
require.Equal(t, time.Time{}, models[0].LimitsContext.From)
})
t.Run("parsing query model with query limits context", func(t *testing.T) {
from := time.Now().Add(-3000 * time.Second)
fullFrom := time.Now().Add(-1 * time.Hour)
to := time.Now()
queryContext := &backend.QueryDataRequest{
Queries: []backend.DataQuery{
{
JSON: []byte(`
{
"expr": "count_over_time({service_name=\"apache\", __stream_shard__=\"2\"}[$__auto])",
"format": "time_series",
"refId": "A",
"limitsContext": {"expr": "count_over_time({service_name=\"apache\"}[$__auto])", "from": ` + strconv.FormatInt(fullFrom.UnixMilli(), 10) + `, "to": ` + strconv.FormatInt(to.UnixMilli(), 10) + `}
}`,
),
TimeRange: backend.TimeRange{
From: from,
To: to,
},
Interval: time.Second * 15,
MaxDataPoints: 200,
},
},
}
models, err := parseQuery(queryContext, true)
require.NoError(t, err)
require.Equal(t, time.Second*15, models[0].Step)
require.Equal(t, `count_over_time({service_name="apache", __stream_shard__="2"}[15s])`, models[0].Expr)
require.Equal(t, `count_over_time({service_name="apache"}[15s])`, models[0].LimitsContext.Expr)
require.Equal(t, to.Truncate(time.Millisecond), models[0].LimitsContext.To)
require.Equal(t, fullFrom.Truncate(time.Millisecond), models[0].LimitsContext.From)
})
t.Run("interpolate variables, range between 1s and 0.5s", func(t *testing.T) {
expr := "go_goroutines $__interval $__interval_ms $__range $__range_s $__range_ms"
queryType := dataquery.LokiQueryTypeRange
pkg/tsdb/loki/types.go
+6
View File
@@ -11,6 +11,11 @@ import (
type QueryType = dataquery.LokiQueryType
type SupportingQueryType = dataquery.SupportingQueryType
type Direction = dataquery.LokiQueryDirection
type LimitsContext struct {
Expr string
From time.Time
To time.Time
}
const (
QueryTypeRange = dataquery.LokiQueryTypeRange
@@ -42,4 +47,5 @@ type lokiQuery struct {
RefID string
SupportingQueryType SupportingQueryType
Scopes []scope.ScopeFilter
LimitsContext LimitsContext
}