github.com/choria-io/go-choria@v0.28.1-0.20240416190746-b3bf9c7d5a45/providers/kv/kv.go (about)

     1  // Copyright (c) 2021-2024, R.I. Pienaar and the Choria Project contributors
     2  //
     3  // SPDX-License-Identifier: Apache-2.0
     4  
     5  package kv
     6  
     7  import (
     8  	"fmt"
     9  	"time"
    10  
    11  	"github.com/nats-io/jsm.go"
    12  	"github.com/nats-io/jsm.go/api"
    13  	"github.com/nats-io/nats.go"
    14  )
    15  
    16  // Option configures a KV bucket
    17  type Option func(*options)
    18  
    19  type options struct {
    20  	name          string
    21  	description   string
    22  	maxValSize    int32
    23  	history       uint8
    24  	ttl           time.Duration
    25  	maxBucketSize int64
    26  	replicas      int
    27  	direct        bool
    28  }
    29  
    30  func WithTTL(ttl time.Duration) Option {
    31  	return func(o *options) { o.ttl = ttl }
    32  }
    33  
    34  func WithHistory(h uint8) Option {
    35  	return func(o *options) { o.history = h }
    36  }
    37  
    38  func WithReplicas(r int) Option {
    39  	return func(o *options) { o.replicas = r }
    40  }
    41  
    42  func WithMaxBucketSize(s int64) Option {
    43  	return func(o *options) { o.maxBucketSize = s }
    44  }
    45  
    46  func WithMaxValueSize(s int32) Option {
    47  	return func(o *options) { o.maxValSize = s }
    48  }
    49  
    50  func WithoutDirectAccess() Option {
    51  	return func(o *options) { o.direct = false }
    52  }
    53  
    54  func DeleteKV(nc *nats.Conn, kv nats.KeyValue) error {
    55  	status, err := kv.Status()
    56  	if err != nil {
    57  		return err
    58  	}
    59  	nfo := status.(*nats.KeyValueBucketStatus).StreamInfo()
    60  
    61  	js, err := nc.JetStream()
    62  	if err != nil {
    63  		return err
    64  	}
    65  
    66  	return js.DeleteStream(nfo.Config.Name)
    67  }
    68  
    69  func LoadKV(nc *nats.Conn, name string) (nats.KeyValue, error) {
    70  	js, err := nc.JetStream()
    71  	if err != nil {
    72  		return nil, fmt.Errorf("failed to connect to Choria Streams: %s", err)
    73  	}
    74  
    75  	return js.KeyValue(name)
    76  }
    77  
    78  func NewKV(nc *nats.Conn, name string, create bool, opts ...Option) (nats.KeyValue, error) {
    79  	opt := &options{
    80  		name:        name,
    81  		replicas:    1,
    82  		direct:      true,
    83  		description: "Choria Streams Key-Value Bucket",
    84  	}
    85  
    86  	for _, o := range opts {
    87  		o(opt)
    88  	}
    89  
    90  	kv, err := LoadKV(nc, name)
    91  	if err == nil {
    92  		return kv, nil
    93  	}
    94  
    95  	if !create {
    96  		return nil, fmt.Errorf("failed to load Choria Key-Value store %s: %s", name, err)
    97  	}
    98  
    99  	mgr, err := jsm.New(nc)
   100  	if err != nil {
   101  		return nil, fmt.Errorf("failed to connect to Choria Streams: %s", err)
   102  	}
   103  
   104  	cfg := api.StreamConfig{
   105  		Name:          fmt.Sprintf("KV_%s", name),
   106  		Subjects:      []string{fmt.Sprintf("$KV.%s.>", name)},
   107  		Retention:     api.LimitsPolicy,
   108  		MaxMsgsPer:    int64(opt.history),
   109  		MaxBytes:      opt.maxBucketSize,
   110  		MaxAge:        opt.ttl,
   111  		Replicas:      opt.replicas,
   112  		AllowDirect:   opt.direct,
   113  		MaxConsumers:  -1,
   114  		MaxMsgs:       -1,
   115  		MaxMsgSize:    -1,
   116  		Storage:       api.FileStorage,
   117  		Discard:       api.DiscardNew,
   118  		Duplicates:    2 * time.Minute,
   119  		RollupAllowed: true,
   120  		DenyDelete:    true,
   121  	}
   122  
   123  	if cfg.Duplicates > cfg.MaxAge {
   124  		cfg.Duplicates = cfg.MaxAge
   125  	}
   126  
   127  	_, err = mgr.NewStreamFromDefault(cfg.Name, cfg)
   128  	if err != nil {
   129  		return nil, fmt.Errorf("failed to create key-value bucket: %s", err)
   130  	}
   131  
   132  	return LoadKV(nc, name)
   133  }