github.com/yankunsam/loki/v2@v2.6.3-0.20220817130409-389df5235c27/pkg/util/marshal/marshal.go (about)

     1  // Package marshal converts internal objects to loghttp model objects.  This
     2  // package is designed to work with models in pkg/loghttp.
     3  package marshal
     4  
     5  import (
     6  	"io"
     7  
     8  	"github.com/gorilla/websocket"
     9  	jsoniter "github.com/json-iterator/go"
    10  
    11  	"github.com/grafana/loki/pkg/loghttp"
    12  	legacy "github.com/grafana/loki/pkg/loghttp/legacy"
    13  	"github.com/grafana/loki/pkg/logproto"
    14  	"github.com/grafana/loki/pkg/logqlmodel"
    15  	"github.com/grafana/loki/pkg/storage/stores/index/stats"
    16  )
    17  
    18  // WriteQueryResponseJSON marshals the promql.Value to v1 loghttp JSON and then
    19  // writes it to the provided io.Writer.
    20  func WriteQueryResponseJSON(v logqlmodel.Result, w io.Writer) error {
    21  	value, err := NewResultValue(v.Data)
    22  	if err != nil {
    23  		return err
    24  	}
    25  
    26  	q := loghttp.QueryResponse{
    27  		Status: "success",
    28  		Data: loghttp.QueryResponseData{
    29  			ResultType: value.Type(),
    30  			Result:     value,
    31  			Statistics: v.Statistics,
    32  		},
    33  	}
    34  
    35  	return jsoniter.NewEncoder(w).Encode(q)
    36  }
    37  
    38  // WriteLabelResponseJSON marshals a logproto.LabelResponse to v1 loghttp JSON
    39  // and then writes it to the provided io.Writer.
    40  func WriteLabelResponseJSON(l logproto.LabelResponse, w io.Writer) error {
    41  	v1Response := loghttp.LabelResponse{
    42  		Status: "success",
    43  		Data:   l.GetValues(),
    44  	}
    45  
    46  	return jsoniter.NewEncoder(w).Encode(v1Response)
    47  }
    48  
    49  // WebsocketWriter knows how to write message to a websocket connection.
    50  type WebsocketWriter interface {
    51  	WriteMessage(int, []byte) error
    52  }
    53  
    54  // WriteTailResponseJSON marshals the legacy.TailResponse to v1 loghttp JSON and
    55  // then writes it to the provided connection.
    56  func WriteTailResponseJSON(r legacy.TailResponse, c WebsocketWriter) error {
    57  	v1Response, err := NewTailResponse(r)
    58  	if err != nil {
    59  		return err
    60  	}
    61  	data, err := jsoniter.Marshal(v1Response)
    62  	if err != nil {
    63  		return err
    64  	}
    65  	return c.WriteMessage(websocket.TextMessage, data)
    66  }
    67  
    68  // WriteSeriesResponseJSON marshals a logproto.SeriesResponse to v1 loghttp JSON and then
    69  // writes it to the provided io.Writer.
    70  func WriteSeriesResponseJSON(r logproto.SeriesResponse, w io.Writer) error {
    71  	adapter := &seriesResponseAdapter{
    72  		Status: "success",
    73  		Data:   make([]map[string]string, 0, len(r.GetSeries())),
    74  	}
    75  
    76  	for _, series := range r.GetSeries() {
    77  		adapter.Data = append(adapter.Data, series.GetLabels())
    78  	}
    79  
    80  	return jsoniter.NewEncoder(w).Encode(adapter)
    81  }
    82  
    83  // This struct exists primarily because we can't specify a repeated map in proto v3.
    84  // Otherwise, we'd use that + gogoproto.jsontag to avoid this layer of indirection
    85  type seriesResponseAdapter struct {
    86  	Status string              `json:"status"`
    87  	Data   []map[string]string `json:"data"`
    88  }
    89  
    90  // WriteIndexStatsResponseJSON marshals a gatewaypb.Stats to JSON and then
    91  // writes it to the provided io.Writer.
    92  func WriteIndexStatsResponseJSON(r *stats.Stats, w io.Writer) error {
    93  	return jsoniter.NewEncoder(w).Encode(r)
    94  }