95 lines
1.7 KiB
Go
95 lines
1.7 KiB
Go
package main
|
|
|
|
import (
|
|
"context"
|
|
"io/ioutil"
|
|
"path"
|
|
"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 NewTestDBIn(d string) BBolt {
|
|
d, err := ioutil.TempDir(d, "test-db-*")
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
db, err := NewDB(path.Join(d, "bb"))
|
|
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)
|
|
})
|
|
}
|