initial commit

pkg/store/boltstore.go

@@ -0,0 +1,140 @@
+// Copyright (C) 2022 Marius Schellenberger
+
+package store
+
+import (
+	"encoding/json"
+	"io"
+
+	bolt "go.etcd.io/bbolt"
+)
+
+const (
+	defaultBoltBucket = "default"
+	fileMode          = 0o600
+)
+
+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()
+}