github.com/grafana/pyroscope@v1.18.0/pkg/settings/recording/store.go (about)

     1  package recording
     2  
     3  import (
     4  	"context"
     5  	"encoding/json"
     6  	"fmt"
     7  
     8  	"connectrpc.com/connect"
     9  	"github.com/go-kit/log"
    10  	"github.com/thanos-io/objstore"
    11  	"google.golang.org/protobuf/encoding/protojson"
    12  
    13  	settingsv1 "github.com/grafana/pyroscope/api/gen/proto/go/settings/v1"
    14  	"github.com/grafana/pyroscope/pkg/settings/store"
    15  )
    16  
    17  func newBucketStore(logger log.Logger, bucket objstore.Bucket, key store.Key) *bucketStore {
    18  	bs := &bucketStore{
    19  		logger: logger,
    20  	}
    21  
    22  	bs.store = store.New(logger, bucket, key, &storeHelper{
    23  		b: bs,
    24  	})
    25  	return bs
    26  }
    27  
    28  type bucketStore struct {
    29  	logger log.Logger
    30  	store  *store.GenericStore[*settingsv1.RecordingRuleStore, *storeHelper]
    31  }
    32  
    33  func (b *bucketStore) Get(ctx context.Context, id string) (*settingsv1.RecordingRuleStore, error) {
    34  	var rule *settingsv1.RecordingRuleStore
    35  	err := b.store.Read(ctx, func(ctx context.Context, c *store.Collection[*settingsv1.RecordingRuleStore]) error {
    36  		for _, r := range c.Elements {
    37  			if r.Id != id {
    38  				continue
    39  			}
    40  
    41  			rule = r
    42  			break
    43  		}
    44  		return nil
    45  	})
    46  	if err != nil {
    47  		return nil, err
    48  	}
    49  	return rule, nil
    50  }
    51  
    52  func (b *bucketStore) List(ctx context.Context) (*settingsv1.RecordingRulesStore, error) {
    53  	var rules *settingsv1.RecordingRulesStore
    54  	err := b.store.Read(ctx, func(ctx context.Context, c *store.Collection[*settingsv1.RecordingRuleStore]) error {
    55  		rules = &settingsv1.RecordingRulesStore{
    56  			Rules:      c.Elements,
    57  			Generation: c.Generation,
    58  		}
    59  		return nil
    60  	})
    61  	if err != nil {
    62  		return nil, err
    63  	}
    64  
    65  	return rules, nil
    66  }
    67  
    68  func (b *bucketStore) Upsert(ctx context.Context, newRule *settingsv1.RecordingRuleStore) (*settingsv1.RecordingRuleStore, error) {
    69  	err := b.store.Upsert(ctx, newRule, &newRule.Generation)
    70  	if err != nil {
    71  		return nil, err
    72  	}
    73  
    74  	return newRule, nil
    75  }
    76  
    77  func (b *bucketStore) Delete(ctx context.Context, ruleID string) error {
    78  	err := b.store.Delete(ctx, ruleID)
    79  	if err == store.ErrElementNotFound {
    80  		return connect.NewError(connect.CodeNotFound, fmt.Errorf("no rule with ID='%s' found", ruleID))
    81  	}
    82  	return nil
    83  }
    84  
    85  type storeHelper struct {
    86  	b *bucketStore
    87  }
    88  
    89  func (*storeHelper) SetGeneration(rule *settingsv1.RecordingRuleStore, generation int64) {
    90  	rule.Generation = generation
    91  }
    92  
    93  func (*storeHelper) GetGeneration(rule *settingsv1.RecordingRuleStore) int64 {
    94  	return rule.Generation
    95  }
    96  
    97  func (*storeHelper) FromStore(storeBytes json.RawMessage) (*settingsv1.RecordingRuleStore, error) {
    98  	var store settingsv1.RecordingRuleStore
    99  	err := protojson.Unmarshal(storeBytes, &store)
   100  	if err != nil {
   101  		return nil, fmt.Errorf("error unmarshaling json from store: %w", err)
   102  	}
   103  
   104  	return &store, nil
   105  }
   106  
   107  func (*storeHelper) ToStore(rule *settingsv1.RecordingRuleStore) (json.RawMessage, error) {
   108  	return protojson.Marshal(rule)
   109  }
   110  
   111  func (*storeHelper) ID(rule *settingsv1.RecordingRuleStore) string {
   112  	return rule.Id
   113  }
   114  
   115  func (*storeHelper) TypePath() string {
   116  	return "settings/recording_rule.v1"
   117  }