This commit is contained in:
bel
2024-04-11 23:01:40 -06:00
parent def095e0e8
commit ea902cef86
7 changed files with 98 additions and 17 deletions

View File

@@ -9,9 +9,15 @@ import (
"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
bkt string
db *bbolt.DB
}
func NewTestDB() BBolt {
@@ -37,9 +43,9 @@ func (bb BBolt) Close() error {
return bb.db.Close()
}
func (bb BBolt) ForEach(ctx context.Context, cb func(string, []byte) error) error {
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(bb.bkt))
bkt := tx.Bucket([]byte(db))
if bkt == nil {
return nil
}
@@ -55,10 +61,10 @@ func (bb BBolt) ForEach(ctx context.Context, cb func(string, []byte) error) erro
})
}
func (bb BBolt) Get(_ context.Context, id string) ([]byte, error) {
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(bb.bkt))
bkt := tx.Bucket([]byte(db))
if bkt == nil {
return nil
}
@@ -69,17 +75,20 @@ func (bb BBolt) Get(_ context.Context, id string) ([]byte, error) {
return b, err
}
func (bb BBolt) Set(_ context.Context, id string, value []byte) error {
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(bb.bkt))
bkt := tx.Bucket([]byte(db))
if bkt == nil {
var err error
bkt, err = tx.CreateBucket([]byte(bb.bkt))
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)
})
}