github.com/grafana/pyroscope@v1.18.0/pkg/metastore/query_service.go (about)

     1  package metastore
     2  
     3  import (
     4  	"context"
     5  	"errors"
     6  	"time"
     7  
     8  	"github.com/go-kit/log"
     9  	"github.com/go-kit/log/level"
    10  	"github.com/opentracing/opentracing-go"
    11  	"github.com/opentracing/opentracing-go/ext"
    12  	otlog "github.com/opentracing/opentracing-go/log"
    13  	"go.etcd.io/bbolt"
    14  	"google.golang.org/grpc/codes"
    15  	"google.golang.org/grpc/status"
    16  
    17  	metastorev1 "github.com/grafana/pyroscope/api/gen/proto/go/metastore/v1"
    18  	typesv1 "github.com/grafana/pyroscope/api/gen/proto/go/types/v1"
    19  	"github.com/grafana/pyroscope/pkg/metastore/index"
    20  	"github.com/grafana/pyroscope/pkg/metastore/raftnode"
    21  )
    22  
    23  type IndexQuerier interface {
    24  	QueryMetadata(*bbolt.Tx, context.Context, index.MetadataQuery) ([]*metastorev1.BlockMeta, error)
    25  	QueryMetadataLabels(*bbolt.Tx, context.Context, index.MetadataQuery) ([]*typesv1.Labels, error)
    26  }
    27  
    28  type QueryService struct {
    29  	metastorev1.MetadataQueryServiceServer
    30  
    31  	logger log.Logger
    32  	state  State
    33  	index  IndexQuerier
    34  }
    35  
    36  func NewQueryService(
    37  	logger log.Logger,
    38  	state State,
    39  	index IndexQuerier,
    40  ) *QueryService {
    41  	return &QueryService{
    42  		logger: logger,
    43  		state:  state,
    44  		index:  index,
    45  	}
    46  }
    47  
    48  func (svc *QueryService) QueryMetadata(
    49  	ctx context.Context,
    50  	req *metastorev1.QueryMetadataRequest,
    51  ) (resp *metastorev1.QueryMetadataResponse, err error) {
    52  	span, ctx := opentracing.StartSpanFromContext(ctx, "QueryService.QueryMetadata")
    53  	defer func() {
    54  		if err != nil {
    55  			ext.LogError(span, err)
    56  		}
    57  		span.Finish()
    58  	}()
    59  	span.SetTag("tenant_id", req.GetTenantId())
    60  	span.SetTag("start_time", req.GetStartTime())
    61  	span.SetTag("end_time", req.GetEndTime())
    62  	span.SetTag("labels", len(req.GetLabels()))
    63  	if q := req.GetQuery(); q != "" {
    64  		span.LogFields(otlog.String("query", q))
    65  	}
    66  
    67  	read := func(tx *bbolt.Tx, _ raftnode.ReadIndex) {
    68  		resp, err = svc.queryMetadata(ctx, tx, req)
    69  	}
    70  	if readErr := svc.state.ConsistentRead(ctx, read); readErr != nil {
    71  		return nil, status.Error(codes.Unavailable, readErr.Error())
    72  	}
    73  	return resp, err
    74  }
    75  
    76  func (svc *QueryService) queryMetadata(
    77  	ctx context.Context,
    78  	tx *bbolt.Tx,
    79  	req *metastorev1.QueryMetadataRequest,
    80  ) (resp *metastorev1.QueryMetadataResponse, err error) {
    81  	span, ctx := opentracing.StartSpanFromContext(ctx, "QueryService.indexQueryMetadata")
    82  	defer func() {
    83  		if err != nil {
    84  			ext.LogError(span, err)
    85  		}
    86  		span.Finish()
    87  	}()
    88  
    89  	metas, err := svc.index.QueryMetadata(tx, ctx, index.MetadataQuery{
    90  		Tenant:    req.TenantId,
    91  		StartTime: time.UnixMilli(req.StartTime),
    92  		EndTime:   time.UnixMilli(req.EndTime),
    93  		Expr:      req.Query,
    94  		Labels:    req.Labels,
    95  	})
    96  	if err == nil {
    97  		span.SetTag("result_count", len(metas))
    98  		return &metastorev1.QueryMetadataResponse{Blocks: metas}, nil
    99  	}
   100  	var invalid *index.InvalidQueryError
   101  	if errors.As(err, &invalid) {
   102  		return nil, status.Error(codes.InvalidArgument, err.Error())
   103  	}
   104  	level.Error(svc.logger).Log("msg", "failed to query metadata", "err", err)
   105  	return nil, status.Error(codes.Internal, err.Error())
   106  }
   107  
   108  func (svc *QueryService) QueryMetadataLabels(
   109  	ctx context.Context,
   110  	req *metastorev1.QueryMetadataLabelsRequest,
   111  ) (resp *metastorev1.QueryMetadataLabelsResponse, err error) {
   112  	span, ctx := opentracing.StartSpanFromContext(ctx, "QueryService.QueryMetadataLabels")
   113  	defer func() {
   114  		if err != nil {
   115  			ext.LogError(span, err)
   116  		}
   117  		span.Finish()
   118  	}()
   119  
   120  	span.SetTag("tenant_id", req.GetTenantId())
   121  	span.SetTag("start_time", req.GetStartTime())
   122  	span.SetTag("end_time", req.GetEndTime())
   123  	span.SetTag("labels", len(req.GetLabels()))
   124  	if q := req.GetQuery(); q != "" {
   125  		span.LogFields(otlog.String("query", q))
   126  	}
   127  
   128  	read := func(tx *bbolt.Tx, _ raftnode.ReadIndex) {
   129  		resp, err = svc.queryMetadataLabels(ctx, tx, req)
   130  	}
   131  	if readErr := svc.state.ConsistentRead(ctx, read); readErr != nil {
   132  		return nil, status.Error(codes.Unavailable, readErr.Error())
   133  	}
   134  	return resp, err
   135  }
   136  
   137  func (svc *QueryService) queryMetadataLabels(
   138  	ctx context.Context,
   139  	tx *bbolt.Tx,
   140  	req *metastorev1.QueryMetadataLabelsRequest,
   141  ) (resp *metastorev1.QueryMetadataLabelsResponse, err error) {
   142  	span, ctx := opentracing.StartSpanFromContext(ctx, "QueryService.indexQueryMetadataLabels")
   143  	defer func() {
   144  		if err != nil {
   145  			ext.LogError(span, err)
   146  		}
   147  		span.Finish()
   148  	}()
   149  
   150  	labels, err := svc.index.QueryMetadataLabels(tx, ctx, index.MetadataQuery{
   151  		Tenant:    req.TenantId,
   152  		StartTime: time.UnixMilli(req.StartTime),
   153  		EndTime:   time.UnixMilli(req.EndTime),
   154  		Expr:      req.Query,
   155  		Labels:    req.Labels,
   156  	})
   157  	if err == nil {
   158  		span.SetTag("result_count", len(labels))
   159  		return &metastorev1.QueryMetadataLabelsResponse{Labels: labels}, nil
   160  	}
   161  	var invalid *index.InvalidQueryError
   162  	if errors.As(err, &invalid) {
   163  		return nil, status.Error(codes.InvalidArgument, err.Error())
   164  	}
   165  	level.Error(svc.logger).Log("msg", "failed to query metadata labels", "err", err)
   166  	return nil, status.Error(codes.Internal, err.Error())
   167  }