feat: add header to skip split by interval

pkg/loki/modules.go

@@ -431,6 +431,7 @@ func (t *Loki) initQuerier() (services.Service, error) {
 	toMerge := []middleware.Interface{
 		httpreq.ExtractQueryMetricsMiddleware(),
 		httpreq.ExtractQueryTagsMiddleware(),
+		httpreq.ExtractQueryNoSplitMiddleware(),
 		httpreq.PropagateHeadersMiddleware(httpreq.LokiEncodingFlagsHeader, httpreq.LokiDisablePipelineWrappersHeader),
 		serverutil.RecoveryHTTPMiddleware,
 		t.HTTPAuthMiddleware,
@@ -1128,6 +1129,7 @@ func (t *Loki) initQueryFrontend() (_ services.Service, err error) {
 
 	toMerge := []middleware.Interface{
 		httpreq.ExtractQueryTagsMiddleware(),
+		httpreq.ExtractQueryNoSplitMiddleware(),
 		httpreq.PropagateHeadersMiddleware(httpreq.LokiActorPathHeader, httpreq.LokiEncodingFlagsHeader, httpreq.LokiDisablePipelineWrappersHeader),
 		serverutil.RecoveryHTTPMiddleware,
 		t.HTTPAuthMiddleware,
@@ -1148,6 +1150,7 @@ func (t *Loki) initQueryFrontend() (_ services.Service, err error) {
 	if t.Cfg.Frontend.TailProxyURL != "" && !t.isModuleActive(Querier) {
 		httpMiddleware := middleware.Merge(
 			httpreq.ExtractQueryTagsMiddleware(),
+		  httpreq.ExtractQueryNoSplitMiddleware(),
 			t.HTTPAuthMiddleware,
 			queryrange.StatsHTTPMiddleware,
 		)

pkg/querier/queryrange/split_by_interval.go

@@ -14,6 +14,7 @@ import (
 	"github.com/prometheus/common/model"
 
 	"github.com/grafana/loki/v3/pkg/util/constants"
+	"github.com/grafana/loki/v3/pkg/util/httpreq"
 	"github.com/grafana/loki/v3/pkg/util/math"
 
 	"github.com/grafana/dskit/tenant"
@@ -182,6 +183,11 @@ func (h *splitByInterval) Do(ctx context.Context, r queryrangebase.Request) (que
 		return nil, httpgrpc.Errorf(http.StatusBadRequest, "%s", err.Error())
 	}
 
+	noSplit := httpreq.ExtractQueryNoSplitFromContext(ctx)
+	if noSplit {
+		return h.next.Do(ctx, r)
+	}
+
 	var interval time.Duration
 	switch r.(type) {
 	case *LokiSeriesRequest, *LabelRequest:

pkg/querier/queryrange/split_by_interval_test.go

@@ -24,6 +24,7 @@ import (
 	"github.com/grafana/loki/v3/pkg/storage/config"
 	"github.com/grafana/loki/v3/pkg/storage/stores/index/seriesvolume"
 	"github.com/grafana/loki/v3/pkg/util"
+	"github.com/grafana/loki/v3/pkg/util/httpreq"
 )
 
 var nilMetrics = NewSplitByMetrics(nil)
@@ -1408,6 +1409,7 @@ func Test_splitByInterval_Do(t *testing.T) {
 		name string
 		req  *LokiRequest
 		want *LokiResponse
+		ctx  context.Context
 	}{
 		{
 			"backward",
@@ -1441,6 +1443,38 @@ func Test_splitByInterval_Do(t *testing.T) {
 					},
 				},
 			},
+			ctx,
+		},
+		{
+			"backward - no split header",
+			&LokiRequest{
+				StartTs:   time.Unix(0, 0),
+				EndTs:     time.Unix(0, (4 * time.Hour).Nanoseconds()),
+				Query:     "",
+				Limit:     1000,
+				Step:      1,
+				Direction: logproto.BACKWARD,
+				Path:      "/api/prom/query_range",
+			},
+			&LokiResponse{
+				Status:     loghttp.QueryStatusSuccess,
+				Direction:  logproto.BACKWARD,
+				Limit:      1000,
+				Version:    1,
+				Statistics: stats.Result{Summary: stats.Summary{Splits: 0}},
+				Data: LokiData{
+					ResultType: loghttp.ResultTypeStream,
+					Result: []logproto.Stream{
+						{
+							Labels: `{foo="bar", level="debug"}`,
+							Entries: []logproto.Entry{
+								{Timestamp: time.Unix(0, 0), Line: fmt.Sprintf("%d", 0)},
+							},
+						},
+					},
+				},
+			},
+			httpreq.InjectQueryNoSplit(ctx, true),
 		},
 		{
 			"forward",
@@ -1474,6 +1508,38 @@ func Test_splitByInterval_Do(t *testing.T) {
 					},
 				},
 			},
+			ctx,
+		},
+		{
+			"forward - no split header",
+			&LokiRequest{
+				StartTs:   time.Unix(0, 0),
+				EndTs:     time.Unix(0, (4 * time.Hour).Nanoseconds()),
+				Query:     "",
+				Limit:     1000,
+				Step:      1,
+				Direction: logproto.FORWARD,
+				Path:      "/api/prom/query_range",
+			},
+			&LokiResponse{
+				Status:     loghttp.QueryStatusSuccess,
+				Direction:  logproto.FORWARD,
+				Statistics: stats.Result{Summary: stats.Summary{Splits: 0}},
+				Limit:      1000,
+				Version:    1,
+				Data: LokiData{
+					ResultType: loghttp.ResultTypeStream,
+					Result: []logproto.Stream{
+						{
+							Labels: `{foo="bar", level="debug"}`,
+							Entries: []logproto.Entry{
+								{Timestamp: time.Unix(0, 0), Line: fmt.Sprintf("%d", 0)},
+							},
+						},
+					},
+				},
+			},
+			httpreq.InjectQueryNoSplit(ctx, true),
 		},
 		{
 			"forward limited",
@@ -1505,6 +1571,7 @@ func Test_splitByInterval_Do(t *testing.T) {
 					},
 				},
 			},
+			ctx,
 		},
 		{
 			"backward limited",
@@ -1536,12 +1603,13 @@ func Test_splitByInterval_Do(t *testing.T) {
 					},
 				},
 			},
+			ctx,
 		},
 	}
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			res, err := split.Do(ctx, tt.req)
+			res, err := split.Do(tt.ctx, tt.req)
 			require.NoError(t, err)
 			require.Equal(t, tt.want, res)
 		})

pkg/util/httpreq/headers.go

@@ -14,6 +14,7 @@ var (
 	// LokiActorPathHeader is the name of the header e.g. used to enqueue requests in hierarchical queues.
 	LokiActorPathHeader               = "X-Loki-Actor-Path"
 	LokiDisablePipelineWrappersHeader = "X-Loki-Disable-Pipeline-Wrappers"
+	QueryNoSplitHTTPHeader            = "X-Query-No-Split"
 
 	// LokiActorPathDelimiter is the delimiter used to serialise the hierarchy of the actor.
 	LokiActorPathDelimiter = "|"
@@ -55,3 +56,31 @@ func InjectActorPath(ctx context.Context, value string) context.Context {
 func InjectHeader(ctx context.Context, key, value string) context.Context {
 	return context.WithValue(ctx, headerContextKey(key), value)
 }
+
+func ExtractQueryNoSplitMiddleware() middleware.Interface {
+	return middleware.Func(func(next http.Handler) http.Handler {
+		return http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
+			ctx := req.Context()
+
+			noSplit := ExtractQueryNoSplitFromHTTP(req)
+			ctx = InjectQueryNoSplit(ctx, noSplit)
+			req = req.WithContext(ctx)
+			next.ServeHTTP(w, req)
+		})
+	})
+}
+
+func ExtractQueryNoSplitFromHTTP(req *http.Request) bool {
+	tags := req.Header.Get(string(QueryNoSplitHTTPHeader))
+	return tags != "" && strings.ToLower(tags) == "true"
+}
+
+func ExtractQueryNoSplitFromContext(ctx context.Context) bool {
+	// if the cast fails then v will be an empty string
+	v, _ := ctx.Value(QueryNoSplitHTTPHeader).(bool)
+	return v
+}
+
+func InjectQueryNoSplit(ctx context.Context, noSplit bool) context.Context {
+	return context.WithValue(ctx, QueryNoSplitHTTPHeader, noSplit)
+}

pkg/util/httpreq/headers_test.go

@@ -0,0 +1,61 @@
+package httpreq
+
+import (
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestQueryNoSplit(t *testing.T) {
+	for _, tc := range []struct {
+		desc  string
+		in    string
+		exp   bool
+		error bool
+	}{
+		{
+			desc: "true",
+			in:   `true`,
+			exp:  true,
+		},
+		{
+			desc: "false",
+			in:   `false`,
+			exp:  false,
+		},
+		{
+			desc: "absent",
+			in:   ``,
+			exp:  false,
+		},
+		{
+			desc: "garbage",
+			in:   `garbage`,
+			exp:  false,
+		},
+	} {
+		t.Run(tc.desc, func(t *testing.T) {
+			req := httptest.NewRequest("GET", "http://testing.com", nil)
+			if tc.in != "" {
+				req.Header.Set(string(QueryNoSplitHTTPHeader), tc.in)
+			}
+
+			w := httptest.NewRecorder()
+			checked := false
+			mware := ExtractQueryNoSplitMiddleware().Wrap(http.HandlerFunc(func(_ http.ResponseWriter, req *http.Request) {
+				noSplit := req.Context().Value(QueryNoSplitHTTPHeader)
+				noSplitBool, ok := noSplit.(bool)
+				require.True(t, ok)
+				require.Equal(t, tc.exp, noSplitBool)
+				checked = true
+			}))
+
+			mware.ServeHTTP(w, req)
+
+			assert.True(t, true, checked)
+		})
+	}
+}