Separate badger store into separate package and switch to outcaste badger

2022-06-08 22:00:41 -07:00
parent a697f3c02a
commit 1ef79b6f00
3 changed files with 99 additions and 31 deletions
--- a/badgerstore/badger.go
+++ b/badgerstore/badger.go
@@ -0,0 +1,166 @@
+package badgerstore
+
+import (
+	"strings"
+
+	"github.com/outcaste-io/badger/v3"
+	"github.com/vmihailenco/msgpack/v5"
+	"go.arsenm.dev/store"
+)
+
+type badgerStore struct {
+	db *badger.DB
+}
+
+// NewBadger creates a new store from a badger database
+func New(db *badger.DB) store.Store {
+	return badgerStore{db}
+}
+
+func (b badgerStore) Set(key string, val interface{}) error {
+	return b.db.Update(func(txn *badger.Txn) error {
+		data, err := msgpack.Marshal(val)
+		if err != nil {
+			return err
+		}
+		return txn.Set([]byte(key), data)
+	})
+}
+
+func (b badgerStore) Get(key string, out interface{}) error {
+	return b.db.View(func(txn *badger.Txn) error {
+		item, err := txn.Get([]byte(key))
+		if err != nil {
+			return err
+		}
+
+		return item.Value(func(val []byte) error {
+			return msgpack.Unmarshal(val, out)
+		})
+	})
+}
+
+func (b badgerStore) GetItem(key string) (*store.Item, error) {
+	out := &store.Item{}
+	err := b.db.View(func(txn *badger.Txn) error {
+		item, err := txn.Get([]byte(key))
+		if err != nil {
+			return err
+		}
+
+		key := string(item.KeyCopy(nil))
+		data, err := item.ValueCopy(nil)
+		if err != nil {
+			return err
+		}
+
+		out.Key = key
+		out.Data = data
+
+		return nil
+	})
+	return out, err
+}
+
+func (b badgerStore) Delete(key string) error {
+	return b.db.Update(func(txn *badger.Txn) error {
+		return txn.Delete([]byte(key))
+	})
+}
+
+func (b badgerStore) Modify(key string, modify store.ModifyFunc) error {
+	item, err := b.GetItem(key)
+	if err != nil {
+		return err
+	}
+
+	newVal, err := modify(item)
+	if err != nil {
+		return err
+	}
+
+	return b.Set(key, newVal)
+}
+
+func (b badgerStore) Iter() store.Iter {
+	out := make(chan *store.Item, 10)
+	go func() {
+		b.db.View(func(txn *badger.Txn) error {
+			opts := badger.DefaultIteratorOptions
+			it := txn.NewIterator(opts)
+			defer it.Close()
+			for it.Rewind(); it.Valid(); it.Next() {
+				item := it.Item()
+				key := string(item.KeyCopy(nil))
+				valData, err := item.ValueCopy(nil)
+				if err != nil {
+					continue
+				}
+
+				out <- &store.Item{Key: key, Data: valData}
+			}
+			close(out)
+			return nil
+		})
+	}()
+	return out
+}
+
+func (b badgerStore) Bucket(name string) store.Bucket {
+	return badgerBucket{b, name}
+}
+
+type badgerBucket struct {
+	store badgerStore
+	name  string
+}
+
+func (b badgerBucket) Name() string {
+	return b.name
+}
+
+func (b badgerBucket) Set(key string, val interface{}) error {
+	return b.store.Set(b.name+"\u001f"+key, val)
+}
+
+func (b badgerBucket) Get(key string, out interface{}) error {
+	return b.store.Get(b.name+"\u001f"+key, out)
+}
+
+func (b badgerBucket) GetItem(key string) (*store.Item, error) {
+	return b.store.GetItem(b.name + "\u001f" + key)
+}
+
+func (b badgerBucket) Delete(key string) error {
+	return b.store.Delete(b.name + "\u001f" + key)
+}
+
+func (b badgerBucket) Modify(key string, modify store.ModifyFunc) error {
+	return b.store.Modify(b.name+"\u001f"+key, modify)
+}
+
+func (b badgerBucket) Iter() store.Iter {
+	out := make(chan *store.Item, 10)
+	go func() {
+		b.store.db.View(func(txn *badger.Txn) error {
+			opts := badger.DefaultIteratorOptions
+			opts.Prefix = []byte(b.name + "\u001f")
+			it := txn.NewIterator(opts)
+			defer it.Close()
+			for it.Rewind(); it.Valid(); it.Next() {
+				item := it.Item()
+				key := string(item.KeyCopy(nil))
+				key = strings.TrimPrefix(key, b.name+"\u001f")
+				valData, err := item.ValueCopy(nil)
+				if err != nil {
+					continue
+				}
+
+				out <- &store.Item{Key: key, Data: valData}
+			}
+			close(out)
+			return nil
+		})
+	}()
+	return out
+}