github.com/yankunsam/loki/v2@v2.6.3-0.20220817130409-389df5235c27/pkg/querier/queryrange/prometheus.go (about)

     1  package queryrange
     2  
     3  import (
     4  	"bytes"
     5  	"context"
     6  	"io/ioutil"
     7  	"net/http"
     8  
     9  	jsoniter "github.com/json-iterator/go"
    10  	"github.com/opentracing/opentracing-go"
    11  	otlog "github.com/opentracing/opentracing-go/log"
    12  	"github.com/prometheus/common/model"
    13  
    14  	"github.com/grafana/loki/pkg/loghttp"
    15  	"github.com/grafana/loki/pkg/logqlmodel/stats"
    16  	"github.com/grafana/loki/pkg/querier/queryrange/queryrangebase"
    17  )
    18  
    19  var (
    20  	jsonStd   = jsoniter.ConfigCompatibleWithStandardLibrary
    21  	extractor = queryrangebase.PrometheusResponseExtractor{}
    22  )
    23  
    24  // PrometheusExtractor implements Extractor interface
    25  type PrometheusExtractor struct{}
    26  
    27  // Extract wraps the original prometheus cache extractor
    28  func (PrometheusExtractor) Extract(start, end int64, from queryrangebase.Response) queryrangebase.Response {
    29  	response := extractor.Extract(start, end, from.(*LokiPromResponse).Response)
    30  	return &LokiPromResponse{
    31  		Response: response.(*queryrangebase.PrometheusResponse),
    32  	}
    33  }
    34  
    35  // ResponseWithoutHeaders wraps the original prometheus caching without headers
    36  func (PrometheusExtractor) ResponseWithoutHeaders(resp queryrangebase.Response) queryrangebase.Response {
    37  	response := extractor.ResponseWithoutHeaders(resp.(*LokiPromResponse).Response)
    38  	return &LokiPromResponse{
    39  		Response: response.(*queryrangebase.PrometheusResponse),
    40  	}
    41  }
    42  
    43  // encode encodes a Prometheus response and injects Loki stats.
    44  func (p *LokiPromResponse) encode(ctx context.Context) (*http.Response, error) {
    45  	sp := opentracing.SpanFromContext(ctx)
    46  	var (
    47  		b   []byte
    48  		err error
    49  	)
    50  	if p.Response.Data.ResultType == loghttp.ResultTypeVector {
    51  		b, err = p.marshalVector()
    52  	} else {
    53  		b, err = p.marshalMatrix()
    54  	}
    55  	if err != nil {
    56  		return nil, err
    57  	}
    58  
    59  	if sp != nil {
    60  		sp.LogFields(otlog.Int("bytes", len(b)))
    61  	}
    62  
    63  	resp := http.Response{
    64  		Header: http.Header{
    65  			"Content-Type": []string{"application/json"},
    66  		},
    67  		Body:       ioutil.NopCloser(bytes.NewBuffer(b)),
    68  		StatusCode: http.StatusOK,
    69  	}
    70  	return &resp, nil
    71  }
    72  
    73  func (p *LokiPromResponse) marshalVector() ([]byte, error) {
    74  	vec := make(loghttp.Vector, len(p.Response.Data.Result))
    75  	for i, v := range p.Response.Data.Result {
    76  		lbs := make(model.LabelSet, len(v.Labels))
    77  		for _, v := range v.Labels {
    78  			lbs[model.LabelName(v.Name)] = model.LabelValue(v.Value)
    79  		}
    80  		vec[i] = model.Sample{
    81  			Metric:    model.Metric(lbs),
    82  			Timestamp: model.Time(v.Samples[0].TimestampMs),
    83  			Value:     model.SampleValue(v.Samples[0].Value),
    84  		}
    85  	}
    86  	return jsonStd.Marshal(struct {
    87  		Status string `json:"status"`
    88  		Data   struct {
    89  			ResultType string         `json:"resultType"`
    90  			Result     loghttp.Vector `json:"result"`
    91  			Statistics stats.Result   `json:"stats,omitempty"`
    92  		} `json:"data,omitempty"`
    93  		ErrorType string `json:"errorType,omitempty"`
    94  		Error     string `json:"error,omitempty"`
    95  	}{
    96  		Error: p.Response.Error,
    97  		Data: struct {
    98  			ResultType string         `json:"resultType"`
    99  			Result     loghttp.Vector `json:"result"`
   100  			Statistics stats.Result   `json:"stats,omitempty"`
   101  		}{
   102  			ResultType: loghttp.ResultTypeVector,
   103  			Result:     vec,
   104  			Statistics: p.Statistics,
   105  		},
   106  		ErrorType: p.Response.ErrorType,
   107  		Status:    p.Response.Status,
   108  	})
   109  }
   110  
   111  func (p *LokiPromResponse) marshalMatrix() ([]byte, error) {
   112  	// embed response and add statistics.
   113  	return jsonStd.Marshal(struct {
   114  		Status string `json:"status"`
   115  		Data   struct {
   116  			queryrangebase.PrometheusData
   117  			Statistics stats.Result `json:"stats,omitempty"`
   118  		} `json:"data,omitempty"`
   119  		ErrorType string `json:"errorType,omitempty"`
   120  		Error     string `json:"error,omitempty"`
   121  	}{
   122  		Error: p.Response.Error,
   123  		Data: struct {
   124  			queryrangebase.PrometheusData
   125  			Statistics stats.Result `json:"stats,omitempty"`
   126  		}{
   127  			PrometheusData: p.Response.Data,
   128  			Statistics:     p.Statistics,
   129  		},
   130  		ErrorType: p.Response.ErrorType,
   131  		Status:    p.Response.Status,
   132  	})
   133  }