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

     1  package astmapper
     2  
     3  import (
     4  	"encoding/json"
     5  
     6  	"github.com/prometheus/prometheus/model/labels"
     7  	"github.com/prometheus/prometheus/promql/parser"
     8  )
     9  
    10  /*
    11  Design:
    12  
    13  The prometheus api package enforces a (*promql.Engine argument), making it infeasible to do lazy AST
    14  evaluation and substitution from within this package.
    15  This leaves the (storage.Queryable) interface as the remaining target for conducting application level sharding.
    16  
    17  The main idea is to analyze the AST and determine which subtrees can be parallelized. With those in hand, the queries may
    18  be remapped into vector or matrix selectors utilizing a reserved label containing the original query. These may then be parallelized in the storage implementation.
    19  */
    20  
    21  const (
    22  	// QueryLabel is a reserved label containing an embedded query
    23  	QueryLabel = "__cortex_queries__"
    24  	// EmbeddedQueriesMetricName is a reserved label (metric name) denoting an embedded query
    25  	EmbeddedQueriesMetricName = "__embedded_queries__"
    26  )
    27  
    28  // EmbeddedQueries is a wrapper type for encoding queries
    29  type EmbeddedQueries struct {
    30  	Concat []string `json:"Concat"`
    31  }
    32  
    33  // JSONCodec is a Codec that uses JSON representations of EmbeddedQueries structs
    34  var JSONCodec jsonCodec
    35  
    36  type jsonCodec struct{}
    37  
    38  func (c jsonCodec) Encode(queries []string) (string, error) {
    39  	embedded := EmbeddedQueries{
    40  		Concat: queries,
    41  	}
    42  	b, err := json.Marshal(embedded)
    43  	return string(b), err
    44  }
    45  
    46  func (c jsonCodec) Decode(encoded string) (queries []string, err error) {
    47  	var embedded EmbeddedQueries
    48  	err = json.Unmarshal([]byte(encoded), &embedded)
    49  	if err != nil {
    50  		return nil, err
    51  	}
    52  
    53  	return embedded.Concat, nil
    54  }
    55  
    56  // VectorSquash reduces an AST into a single vector query which can be hijacked by a Queryable impl.
    57  // It always uses a VectorSelector as the substitution node.
    58  // This is important because logical/set binops can only be applied against vectors and not matrices.
    59  func VectorSquasher(nodes ...parser.Node) (parser.Expr, error) {
    60  
    61  	// concat OR legs
    62  	strs := make([]string, 0, len(nodes))
    63  	for _, node := range nodes {
    64  		strs = append(strs, node.String())
    65  	}
    66  
    67  	encoded, err := JSONCodec.Encode(strs)
    68  	if err != nil {
    69  		return nil, err
    70  	}
    71  
    72  	embeddedQuery, err := labels.NewMatcher(labels.MatchEqual, QueryLabel, encoded)
    73  	if err != nil {
    74  		return nil, err
    75  	}
    76  
    77  	return &parser.VectorSelector{
    78  		Name:          EmbeddedQueriesMetricName,
    79  		LabelMatchers: []*labels.Matcher{embeddedQuery},
    80  	}, nil
    81  
    82  }