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

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