86 lines
1.5 KiB
Go
86 lines
1.5 KiB
Go
package main
|
|
|
|
import (
|
|
"context"
|
|
"io/ioutil"
|
|
"os"
|
|
"time"
|
|
|
|
"go.etcd.io/bbolt"
|
|
)
|
|
|
|
type BBolt struct {
|
|
db *bbolt.DB
|
|
bkt string
|
|
}
|
|
|
|
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, cb func(string, []byte) error) error {
|
|
return bb.db.View(func(tx *bbolt.Tx) error {
|
|
bkt := tx.Bucket([]byte(bb.bkt))
|
|
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, id string) ([]byte, error) {
|
|
var b []byte
|
|
err := bb.db.View(func(tx *bbolt.Tx) error {
|
|
bkt := tx.Bucket([]byte(bb.bkt))
|
|
if bkt == nil {
|
|
return nil
|
|
}
|
|
|
|
b = bkt.Get([]byte(id))
|
|
return nil
|
|
})
|
|
return b, err
|
|
}
|
|
|
|
func (bb BBolt) Set(_ context.Context, id string, value []byte) error {
|
|
return bb.db.Update(func(tx *bbolt.Tx) error {
|
|
bkt := tx.Bucket([]byte(bb.bkt))
|
|
if bkt == nil {
|
|
var err error
|
|
bkt, err = tx.CreateBucket([]byte(bb.bkt))
|
|
if err != nil {
|
|
return err
|
|
}
|
|
}
|
|
|
|
return bkt.Put([]byte(id), value)
|
|
})
|
|
}
|