github.com/grafana/pyroscope@v1.18.0/pkg/querybackend/query_label_values.go (about)

     1  package querybackend
     2  
     3  import (
     4  	"errors"
     5  	"sort"
     6  	"sync"
     7  
     8  	"github.com/prometheus/prometheus/model/labels"
     9  	"github.com/prometheus/prometheus/storage"
    10  
    11  	queryv1 "github.com/grafana/pyroscope/api/gen/proto/go/query/v1"
    12  	"github.com/grafana/pyroscope/pkg/block"
    13  	"github.com/grafana/pyroscope/pkg/model"
    14  	"github.com/grafana/pyroscope/pkg/phlaredb"
    15  )
    16  
    17  func init() {
    18  	registerQueryType(
    19  		queryv1.QueryType_QUERY_LABEL_VALUES,
    20  		queryv1.ReportType_REPORT_LABEL_VALUES,
    21  		queryLabelValues,
    22  		newLabelValueAggregator,
    23  		false,
    24  		[]block.Section{block.SectionTSDB}...,
    25  	)
    26  }
    27  
    28  func queryLabelValues(q *queryContext, query *queryv1.Query) (*queryv1.Report, error) {
    29  	var values []string
    30  	var err error
    31  	if len(q.req.matchers) == 0 {
    32  		values, err = q.ds.Index().LabelValues(query.LabelValues.LabelName)
    33  	} else {
    34  		values, err = labelValuesForMatchers(q.ds.Index(), query.LabelValues.LabelName, q.req.matchers)
    35  	}
    36  	if err != nil {
    37  		return nil, err
    38  	}
    39  	resp := &queryv1.Report{
    40  		LabelValues: &queryv1.LabelValuesReport{
    41  			Query:       query.LabelValues.CloneVT(),
    42  			LabelValues: values,
    43  		},
    44  	}
    45  	return resp, nil
    46  }
    47  
    48  func labelValuesForMatchers(reader phlaredb.IndexReader, name string, matchers []*labels.Matcher) ([]string, error) {
    49  	postings, err := phlaredb.PostingsForMatchers(reader, nil, matchers...)
    50  	if err != nil {
    51  		return nil, err
    52  	}
    53  	l := make(map[string]struct{})
    54  	for postings.Next() {
    55  		var v string
    56  		if v, err = reader.LabelValueFor(postings.At(), name); err != nil {
    57  			if errors.Is(err, storage.ErrNotFound) {
    58  				continue
    59  			}
    60  			return nil, err
    61  		}
    62  		l[v] = struct{}{}
    63  	}
    64  	if err = postings.Err(); err != nil {
    65  		return nil, err
    66  	}
    67  	values := make([]string, len(l))
    68  	var i int
    69  	for v := range l {
    70  		values[i] = v
    71  		i++
    72  	}
    73  	sort.Strings(values)
    74  	return values, nil
    75  }
    76  
    77  type labelValueAggregator struct {
    78  	init   sync.Once
    79  	query  *queryv1.LabelValuesQuery
    80  	values *model.LabelMerger
    81  }
    82  
    83  func newLabelValueAggregator(*queryv1.InvokeRequest) aggregator {
    84  	return new(labelValueAggregator)
    85  }
    86  
    87  func (m *labelValueAggregator) aggregate(report *queryv1.Report) error {
    88  	r := report.LabelValues
    89  	m.init.Do(func() {
    90  		m.query = r.Query.CloneVT()
    91  		m.values = model.NewLabelMerger()
    92  	})
    93  	m.values.MergeLabelValues(r.LabelValues)
    94  	return nil
    95  }
    96  
    97  func (m *labelValueAggregator) build() *queryv1.Report {
    98  	return &queryv1.Report{
    99  		LabelValues: &queryv1.LabelValuesReport{
   100  			Query:       m.query,
   101  			LabelValues: m.values.LabelValues(),
   102  		},
   103  	}
   104  }