feat: Enable warnings in Loki query responses (#12425)

pkg/ingester/ingester.go

@@ -12,6 +12,8 @@ import (
 	"sync"
 	"time"
 
+	"github.com/grafana/loki/v3/pkg/logqlmodel/metadata"
+
 	lokilog "github.com/grafana/loki/v3/pkg/logql/log"
 
 	"github.com/go-kit/log"
@@ -874,6 +876,7 @@ func (i *Ingester) GetOrCreateInstance(instanceID string) (*instance, error) { /
 func (i *Ingester) Query(req *logproto.QueryRequest, queryServer logproto.Querier_QueryServer) error {
 	// initialize stats collection for ingester queries.
 	_, ctx := stats.NewContext(queryServer.Context())
+	_, ctx = metadata.NewContext(ctx)
 
 	if req.Plan == nil {
 		parsed, err := syntax.ParseLogSelector(req.Selector, true)
@@ -933,6 +936,7 @@ func (i *Ingester) Query(req *logproto.QueryRequest, queryServer logproto.Querie
 func (i *Ingester) QuerySample(req *logproto.SampleQueryRequest, queryServer logproto.Querier_QuerySampleServer) error {
 	// initialize stats collection for ingester queries.
 	_, ctx := stats.NewContext(queryServer.Context())
+	_, ctx = metadata.NewContext(ctx)
 	sp := opentracing.SpanFromContext(ctx)
 
 	// If the plan is empty we want all series to be returned.

pkg/ingester/instance.go

@@ -10,6 +10,8 @@ import (
 	"syscall"
 	"time"
 
+	"github.com/grafana/loki/v3/pkg/logqlmodel/metadata"
+
 	"github.com/grafana/loki/v3/pkg/util/httpreq"
 
 	"github.com/go-kit/log/level"
@@ -981,6 +983,7 @@ type QuerierQueryServer interface {
 
 func sendBatches(ctx context.Context, i iter.EntryIterator, queryServer QuerierQueryServer, limit int32) error {
 	stats := stats.FromContext(ctx)
+	metadata := metadata.FromContext(ctx)
 
 	// send until the limit is reached.
 	for limit != 0 && !isDone(ctx) {
@@ -999,6 +1002,7 @@ func sendBatches(ctx context.Context, i iter.EntryIterator, queryServer QuerierQ
 
 		stats.AddIngesterBatch(int64(batchSize))
 		batch.Stats = stats.Ingester()
+		batch.Warnings = metadata.Warnings()
 
 		if isDone(ctx) {
 			break
@@ -1013,6 +1017,7 @@ func sendBatches(ctx context.Context, i iter.EntryIterator, queryServer QuerierQ
 		}
 
 		stats.Reset()
+		metadata.Reset()
 	}
 	return nil
 }
@@ -1021,6 +1026,7 @@ func sendSampleBatches(ctx context.Context, it iter.SampleIterator, queryServer
 	sp := opentracing.SpanFromContext(ctx)
 
 	stats := stats.FromContext(ctx)
+	metadata := metadata.FromContext(ctx)
 	for !isDone(ctx) {
 		batch, size, err := iter.ReadSampleBatch(it, queryBatchSampleSize)
 		if err != nil {
@@ -1029,6 +1035,8 @@ func sendSampleBatches(ctx context.Context, it iter.SampleIterator, queryServer
 
 		stats.AddIngesterBatch(int64(size))
 		batch.Stats = stats.Ingester()
+		batch.Warnings = metadata.Warnings()
+
 		if isDone(ctx) {
 			break
 		}
@@ -1042,6 +1050,8 @@ func sendSampleBatches(ctx context.Context, it iter.SampleIterator, queryServer
 		}
 
 		stats.Reset()
+		metadata.Reset()
+
 		if sp != nil {
 			sp.LogKV("event", "sent batch", "size", size)
 		}

pkg/iter/entry_iterator.go

@@ -7,6 +7,8 @@ import (
 	"sync"
 	"time"
 
+	"github.com/grafana/loki/v3/pkg/logqlmodel/metadata"
+
 	"github.com/grafana/loki/v3/pkg/logproto"
 	"github.com/grafana/loki/v3/pkg/logqlmodel/stats"
 	"github.com/grafana/loki/v3/pkg/util"
@@ -379,6 +381,7 @@ func (i *queryClientIterator) Next() bool {
 			return false
 		}
 		stats.JoinIngesters(ctx, batch.Stats)
+		_ = metadata.AddWarnings(ctx, batch.Warnings...)
 		i.curr = NewQueryResponseIterator(batch, i.direction)
 	}
 

pkg/iter/sample_iterator.go

@@ -6,6 +6,8 @@ import (
 	"io"
 	"sync"
 
+	"github.com/grafana/loki/v3/pkg/logqlmodel/metadata"
+
 	"github.com/grafana/loki/v3/pkg/logproto"
 	"github.com/grafana/loki/v3/pkg/logqlmodel/stats"
 	"github.com/grafana/loki/v3/pkg/util"
@@ -490,6 +492,8 @@ func (i *sampleQueryClientIterator) Next() bool {
 			return false
 		}
 		stats.JoinIngesters(ctx, batch.Stats)
+		_ = metadata.AddWarnings(ctx, batch.Warnings...)
+
 		i.curr = NewSampleQueryResponseIterator(batch)
 	}
 	return true

pkg/loghttp/query.go

@@ -38,19 +38,34 @@ type QueryStatus string
 const (
 	QueryStatusSuccess = "success"
 	QueryStatusFail    = "fail"
+	// How much stack space to allocate for unescaping JSON strings; if a string longer
+	// than this needs to be escaped, it will result in a heap allocation
+	unescapeStackBufSize = 64
 )
 
 // QueryResponse represents the http json response to a Loki range and instant query
 type QueryResponse struct {
-	Status string            `json:"status"`
-	Data   QueryResponseData `json:"data"`
+	Status   string            `json:"status"`
+	Warnings []string          `json:"warnings,omitempty"`
+	Data     QueryResponseData `json:"data"`
 }
 
 func (q *QueryResponse) UnmarshalJSON(data []byte) error {
 	return jsonparser.ObjectEach(data, func(key, value []byte, dataType jsonparser.ValueType, offset int) error {
 		switch string(key) {
 		case "status":
 			q.Status = string(value)
+		case "warnings":
+			var warnings []string
+			if _, err := jsonparser.ArrayEach(value, func(value []byte, dataType jsonparser.ValueType, offset int, err error) {
+				if dataType == jsonparser.String {
+					warnings = append(warnings, unescapeJSONString(value))
+				}
+			}); err != nil {
+				return err
+			}
+
+			q.Warnings = warnings
 		case "data":
 			var responseData QueryResponseData
 			if err := responseData.UnmarshalJSON(value); err != nil {
@@ -62,6 +77,16 @@ func (q *QueryResponse) UnmarshalJSON(data []byte) error {
 	})
 }
 
+func unescapeJSONString(b []byte) string {
+	var stackbuf [unescapeStackBufSize]byte // stack-allocated array for allocation-free unescaping of small strings
+	bU, err := jsonparser.Unescape(b, stackbuf[:])
+	if err != nil {
+		return ""
+	}
+
+	return string(bU)
+}
+
 // PushRequest models a log stream push but is unmarshalled to proto push format.
 type PushRequest struct {
 	Streams []LogProtoStream `json:"streams"`