122 lines
2.5 KiB
Go
122 lines
2.5 KiB
Go
package store
|
|
|
|
import (
|
|
"strings"
|
|
|
|
"github.com/dgraph-io/badger/v3"
|
|
"github.com/vmihailenco/msgpack/v5"
|
|
)
|
|
|
|
type badgerStore struct {
|
|
db *badger.DB
|
|
}
|
|
|
|
// NewBadger creates a new store from a badger database
|
|
func NewBadger(db *badger.DB) 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) Delete(key string) error {
|
|
return b.db.Update(func(txn *badger.Txn) error {
|
|
return txn.Delete([]byte(key))
|
|
})
|
|
}
|
|
|
|
func (b badgerStore) Iter() Iter {
|
|
out := make(chan *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 <- &Item{Key: key, Data: valData}
|
|
}
|
|
close(out)
|
|
return nil
|
|
})
|
|
}()
|
|
return out
|
|
}
|
|
|
|
func (b badgerStore) Bucket(name string) 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) Delete(key string) error {
|
|
return b.store.Delete(b.name + "\u001f" + key)
|
|
}
|
|
|
|
func (b badgerBucket) Iter() Iter {
|
|
out := make(chan *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 <- &Item{Key: key, Data: valData}
|
|
}
|
|
close(out)
|
|
return nil
|
|
})
|
|
}()
|
|
return out
|
|
}
|