some type stubs

driver.go

@@ -0,0 +1,94 @@
+package main
+
+import (
+	"context"
+	"io/ioutil"
+	"os"
+	"time"
+
+	"go.etcd.io/bbolt"
+)
+
+type Driver interface {
+	Close() error
+	ForEach(context.Context, string, func(string, []byte) error) error
+	Get(context.Context, string, string) ([]byte, error)
+	Set(context.Context, string, string, []byte) error
+}
+
+type BBolt struct {
+	db *bbolt.DB
+}
+
+func NewTestDB() BBolt {
+	d, err := ioutil.TempDir(os.TempDir(), "test-db-*")
+	if err != nil {
+		panic(err)
+	}
+	db, err := NewDB(d)
+	if err != nil {
+		panic(err)
+	}
+	return db
+}
+
+func NewDB(p string) (BBolt, error) {
+	db, err := bbolt.Open(p, 0600, &bbolt.Options{
+		Timeout: time.Second,
+	})
+	return BBolt{db: db}, err
+}
+
+func (bb BBolt) Close() error {
+	return bb.db.Close()
+}
+
+func (bb BBolt) ForEach(ctx context.Context, db string, cb func(string, []byte) error) error {
+	return bb.db.View(func(tx *bbolt.Tx) error {
+		bkt := tx.Bucket([]byte(db))
+		if bkt == nil {
+			return nil
+		}
+
+		c := bkt.Cursor()
+		for k, v := c.First(); k != nil && ctx.Err() == nil; k, v = c.Next() {
+			if err := cb(string(k), v); err != nil {
+				return err
+			}
+		}
+
+		return ctx.Err()
+	})
+}
+
+func (bb BBolt) Get(_ context.Context, db, id string) ([]byte, error) {
+	var b []byte
+	err := bb.db.View(func(tx *bbolt.Tx) error {
+		bkt := tx.Bucket([]byte(db))
+		if bkt == nil {
+			return nil
+		}
+
+		b = bkt.Get([]byte(id))
+		return nil
+	})
+	return b, err
+}
+
+func (bb BBolt) Set(_ context.Context, db, id string, value []byte) error {
+	return bb.db.Update(func(tx *bbolt.Tx) error {
+		bkt := tx.Bucket([]byte(db))
+		if bkt == nil {
+			var err error
+			bkt, err = tx.CreateBucket([]byte(db))
+			if err != nil {
+				return err
+			}
+		}
+
+		if value == nil {
+			return bkt.Delete([]byte(id))
+		}
+		return bkt.Put([]byte(id), value)
+	})
+}