139 lines
2.8 KiB
Go
139 lines
2.8 KiB
Go
// Copyright (C) 2019 Marius Schellenberger
|
|
|
|
package store
|
|
|
|
import (
|
|
"encoding/json"
|
|
"github.com/boltdb/bolt"
|
|
"io"
|
|
)
|
|
|
|
const (
|
|
defaultBoltBucket = "default"
|
|
fileMode = 0640
|
|
)
|
|
|
|
type BoltStore struct {
|
|
Marshaler
|
|
db *bolt.DB
|
|
}
|
|
|
|
func NewBoltStore(file string, m Marshaler) (Store, error) {
|
|
if m == nil {
|
|
m = NewGOB()
|
|
}
|
|
db, err := bolt.Open(file, fileMode, nil)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
if err = db.Update(func(tx *bolt.Tx) (err error) {
|
|
_, err = tx.CreateBucketIfNotExists([]byte(defaultBoltBucket))
|
|
return
|
|
}); err != nil {
|
|
return nil, err
|
|
}
|
|
return &BoltStore{m, db}, nil
|
|
}
|
|
|
|
type dump struct {
|
|
K string `json:"k"`
|
|
V []byte `json:"v"`
|
|
}
|
|
|
|
func (bs *BoltStore) Dump(w io.Writer) (err error) {
|
|
if w == nil {
|
|
return ErrWriterIsNil
|
|
}
|
|
var d []dump
|
|
err = bs.ForEach(func(k string, v []byte) error {
|
|
b := make([]byte, len(v))
|
|
copy(b, v)
|
|
d = append(d, dump{k, b})
|
|
return nil
|
|
})
|
|
if err != nil {
|
|
return
|
|
}
|
|
return json.NewEncoder(w).Encode(d)
|
|
}
|
|
|
|
func (bs *BoltStore) Restore(r io.Reader) (err error) {
|
|
if r == nil {
|
|
return ErrReaderIsNil
|
|
}
|
|
var d []dump
|
|
err = json.NewDecoder(r).Decode(&d)
|
|
if err != nil {
|
|
return
|
|
}
|
|
if err = bs.db.Update(func(tx *bolt.Tx) (err error) {
|
|
if err = tx.DeleteBucket([]byte(defaultBoltBucket)); err != nil {
|
|
return
|
|
}
|
|
_, err = tx.CreateBucket([]byte(defaultBoltBucket))
|
|
return
|
|
}); err != nil {
|
|
return
|
|
}
|
|
return bs.db.Update(func(tx *bolt.Tx) (err error) {
|
|
b := tx.Bucket([]byte(defaultBoltBucket))
|
|
for _, v := range d {
|
|
b.Put([]byte(v.K), v.V)
|
|
}
|
|
return nil
|
|
})
|
|
}
|
|
|
|
func (bs *BoltStore) Get(key string, v interface{}) (err error) {
|
|
err = bs.db.View(func(tx *bolt.Tx) error {
|
|
b := tx.Bucket([]byte(defaultBoltBucket)).Get([]byte(key))
|
|
if b == nil {
|
|
return ErrKeyNotFound
|
|
}
|
|
if v != nil {
|
|
return bs.Unmarshal(b, v)
|
|
}
|
|
return nil
|
|
})
|
|
return
|
|
}
|
|
|
|
func (bs *BoltStore) Set(key string, v interface{}) error {
|
|
return bs.db.Update(func(tx *bolt.Tx) (err error) {
|
|
var b []byte
|
|
if v != nil {
|
|
b, err = bs.Marshal(v)
|
|
if err != nil {
|
|
return
|
|
}
|
|
}
|
|
return tx.Bucket([]byte(defaultBoltBucket)).Put([]byte(key), b)
|
|
})
|
|
}
|
|
|
|
func (bs *BoltStore) ForEach(f func(string, []byte) error) error {
|
|
return bs.db.View(func(tx *bolt.Tx) error {
|
|
return tx.Bucket([]byte(defaultBoltBucket)).ForEach(func(k, v []byte) error {
|
|
return f(string(k), v)
|
|
})
|
|
})
|
|
}
|
|
|
|
func (bs *BoltStore) ForEachPrefix(prefix string, f func(string, []byte) error) error {
|
|
return bs.ForEach(func(k string, v []byte) error {
|
|
if trim, ok := hasTrimPrefix(k, prefix); ok {
|
|
return f(trim, v)
|
|
}
|
|
return nil
|
|
})
|
|
}
|
|
|
|
func (bs *BoltStore) Delete(key string) error {
|
|
return bs.db.Update(func(tx *bolt.Tx) error {
|
|
return tx.Bucket([]byte(defaultBoltBucket)).Delete([]byte(key))
|
|
})
|
|
}
|
|
|
|
func (bs *BoltStore) Close() error {
|
|
return bs.db.Close()
|
|
}
|