// 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() }