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

     1  package queryrangebase
     2  
     3  import (
     4  	"github.com/pkg/errors"
     5  	"github.com/prometheus/common/model"
     6  	"github.com/prometheus/prometheus/model/labels"
     7  	"github.com/prometheus/prometheus/promql"
     8  	"github.com/prometheus/prometheus/promql/parser"
     9  	"github.com/prometheus/prometheus/storage"
    10  
    11  	"github.com/grafana/loki/pkg/logproto"
    12  	"github.com/grafana/loki/pkg/querier/series"
    13  )
    14  
    15  // FromResult transforms a promql query result into a samplestream
    16  func FromResult(res *promql.Result) ([]SampleStream, error) {
    17  	if res.Err != nil {
    18  		// The error could be wrapped by the PromQL engine. We get the error's cause in order to
    19  		// correctly parse the error in parent callers (eg. gRPC response status code extraction).
    20  		return nil, errors.Cause(res.Err)
    21  	}
    22  	switch v := res.Value.(type) {
    23  	case promql.Scalar:
    24  		return []SampleStream{
    25  			{
    26  				Samples: []logproto.LegacySample{
    27  					{
    28  						Value:       v.V,
    29  						TimestampMs: v.T,
    30  					},
    31  				},
    32  			},
    33  		}, nil
    34  
    35  	case promql.Vector:
    36  		res := make([]SampleStream, 0, len(v))
    37  		for _, sample := range v {
    38  			res = append(res, SampleStream{
    39  				Labels:  mapLabels(sample.Metric),
    40  				Samples: mapPoints(sample.Point),
    41  			})
    42  		}
    43  		return res, nil
    44  
    45  	case promql.Matrix:
    46  		res := make([]SampleStream, 0, len(v))
    47  		for _, series := range v {
    48  			res = append(res, SampleStream{
    49  				Labels:  mapLabels(series.Metric),
    50  				Samples: mapPoints(series.Points...),
    51  			})
    52  		}
    53  		return res, nil
    54  
    55  	}
    56  
    57  	return nil, errors.Errorf("Unexpected value type: [%s]", res.Value.Type())
    58  }
    59  
    60  func mapLabels(ls labels.Labels) []logproto.LabelAdapter {
    61  	result := make([]logproto.LabelAdapter, 0, len(ls))
    62  	for _, l := range ls {
    63  		result = append(result, logproto.LabelAdapter(l))
    64  	}
    65  
    66  	return result
    67  }
    68  
    69  func mapPoints(pts ...promql.Point) []logproto.LegacySample {
    70  	result := make([]logproto.LegacySample, 0, len(pts))
    71  
    72  	for _, pt := range pts {
    73  		result = append(result, logproto.LegacySample{
    74  			Value:       pt.V,
    75  			TimestampMs: pt.T,
    76  		})
    77  	}
    78  
    79  	return result
    80  }
    81  
    82  // ResponseToSamples is needed to map back from api response to the underlying series data
    83  func ResponseToSamples(resp Response) ([]SampleStream, error) {
    84  	promRes, ok := resp.(*PrometheusResponse)
    85  	if !ok {
    86  		return nil, errors.Errorf("error invalid response type: %T, expected: %T", resp, &PrometheusResponse{})
    87  	}
    88  	if promRes.Error != "" {
    89  		return nil, errors.New(promRes.Error)
    90  	}
    91  	switch promRes.Data.ResultType {
    92  	case string(parser.ValueTypeVector), string(parser.ValueTypeMatrix):
    93  		return promRes.Data.Result, nil
    94  	}
    95  
    96  	return nil, errors.Errorf(
    97  		"Invalid promql.Value type: [%s]. Only %s and %s supported",
    98  		promRes.Data.ResultType,
    99  		parser.ValueTypeVector,
   100  		parser.ValueTypeMatrix,
   101  	)
   102  }
   103  
   104  // NewSeriesSet returns an in memory storage.SeriesSet from a []SampleStream
   105  // As NewSeriesSet uses NewConcreteSeriesSet to implement SeriesSet, result will be sorted by label names.
   106  func NewSeriesSet(results []SampleStream) storage.SeriesSet {
   107  	set := make([]storage.Series, 0, len(results))
   108  
   109  	for _, stream := range results {
   110  		samples := make([]model.SamplePair, 0, len(stream.Samples))
   111  		for _, sample := range stream.Samples {
   112  			samples = append(samples, model.SamplePair{
   113  				Timestamp: model.Time(sample.TimestampMs),
   114  				Value:     model.SampleValue(sample.Value),
   115  			})
   116  		}
   117  
   118  		ls := make([]labels.Label, 0, len(stream.Labels))
   119  		for _, l := range stream.Labels {
   120  			ls = append(ls, labels.Label(l))
   121  		}
   122  		set = append(set, series.NewConcreteSeries(ls, samples))
   123  	}
   124  	return series.NewConcreteSeriesSet(set)
   125  }