92 lines
1.6 KiB
Go
92 lines
1.6 KiB
Go
package storage
|
|
|
|
import (
|
|
"encoding/json"
|
|
"github.com/dgraph-io/badger/v4"
|
|
"github.com/spf13/viper"
|
|
"time"
|
|
)
|
|
|
|
type BadgerHandler struct {
|
|
db *badger.DB
|
|
}
|
|
|
|
func NewBadgerHandler(cfg *viper.Viper) (*BadgerHandler, error) {
|
|
dbPath := cfg.GetString("storage.kv.badger.dbPath")
|
|
|
|
opts := badger.DefaultOptions(dbPath)
|
|
db, err := badger.Open(opts)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
return &BadgerHandler{
|
|
db: db,
|
|
}, nil
|
|
}
|
|
|
|
func (bh *BadgerHandler) Put(k string, v any) error {
|
|
err := bh.db.Update(func(txn *badger.Txn) error {
|
|
bytes, err := json.Marshal(v)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
return txn.Set([]byte(k), bytes)
|
|
})
|
|
|
|
return err
|
|
}
|
|
|
|
func (bh *BadgerHandler) PutWithTTL(k string, v any, duration time.Duration) error {
|
|
err := bh.db.Update(func(txn *badger.Txn) error {
|
|
bytes, err := json.Marshal(v)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
entry := badger.NewEntry([]byte(k), bytes).WithTTL(duration)
|
|
return txn.SetEntry(entry)
|
|
})
|
|
|
|
return err
|
|
}
|
|
|
|
func (bh *BadgerHandler) Get(k string) (any, error) {
|
|
var value []byte
|
|
|
|
err := bh.db.View(func(txn *badger.Txn) error {
|
|
item, err := txn.Get([]byte(k))
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
err = item.Value(func(val []byte) error {
|
|
value = make([]byte, len(val))
|
|
copy(value, val)
|
|
return nil
|
|
})
|
|
|
|
return err
|
|
})
|
|
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
var result map[string]interface{}
|
|
err = json.Unmarshal(value, &result)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
return result, nil
|
|
}
|
|
|
|
func (bh *BadgerHandler) Delete(k string) error {
|
|
err := bh.db.Update(func(txn *badger.Txn) error {
|
|
err := txn.Delete([]byte(k))
|
|
return err
|
|
})
|
|
|
|
return err
|
|
}
|