139 lines
3.5 KiB
Go
139 lines
3.5 KiB
Go
package universe_badger
|
|
|
|
import (
|
|
"context"
|
|
"github.com/dgraph-io/badger/v4"
|
|
"time"
|
|
"universe.mrlokop.dev/core/errors"
|
|
kv "universe.mrlokop.dev/key_value"
|
|
)
|
|
|
|
func wrapErr(err error) error {
|
|
if errors.Is(err, badger.ErrKeyNotFound) {
|
|
return kv.ErrNoSuchKey.CreateWithCause(err)
|
|
}
|
|
return err
|
|
}
|
|
|
|
type keyValueStorage struct {
|
|
db *badger.DB
|
|
prefix string
|
|
}
|
|
|
|
func (k *keyValueStorage) Get(ctx context.Context, key string) ([]byte, error) {
|
|
var result []byte
|
|
err := k.db.View(func(txn *badger.Txn) error {
|
|
keyBytes := []byte(k.prefix + key)
|
|
item, err := txn.Get(keyBytes)
|
|
if err != nil {
|
|
return wrapErr(err)
|
|
}
|
|
result = make([]byte, item.ValueSize())
|
|
_, err = item.ValueCopy(result)
|
|
return wrapErr(err)
|
|
})
|
|
return result, err
|
|
}
|
|
|
|
func (k *keyValueStorage) Has(ctx context.Context, key string) (bool, error) {
|
|
var exists bool
|
|
err := k.db.View(func(txn *badger.Txn) error {
|
|
keyBytes := []byte(k.prefix + key)
|
|
_, err := txn.Get(keyBytes)
|
|
if err != nil {
|
|
if errors.Is(err, badger.ErrKeyNotFound) {
|
|
exists = false
|
|
return nil
|
|
}
|
|
return wrapErr(err)
|
|
}
|
|
exists = true
|
|
return nil
|
|
})
|
|
return exists, err
|
|
}
|
|
|
|
func (k *keyValueStorage) Set(ctx context.Context, key string, value []byte) error {
|
|
return k.db.Update(func(txn *badger.Txn) error {
|
|
keyBytes := []byte(k.prefix + key)
|
|
return wrapErr(txn.Set(keyBytes, value))
|
|
})
|
|
}
|
|
|
|
func (k *keyValueStorage) SetWithExpiration(ctx context.Context, key string, value []byte, expiresAt time.Time) error {
|
|
return k.db.Update(func(txn *badger.Txn) error {
|
|
keyBytes := []byte(k.prefix + key)
|
|
entry := badger.NewEntry(keyBytes, value)
|
|
entry.ExpiresAt = uint64(expiresAt.Unix())
|
|
return wrapErr(txn.SetEntry(entry))
|
|
})
|
|
}
|
|
|
|
func (k *keyValueStorage) SetIfNotExist(ctx context.Context, key string, value []byte) (bool, error) {
|
|
var wasSet bool
|
|
err := k.db.Update(func(txn *badger.Txn) error {
|
|
keyBytes := []byte(k.prefix + key)
|
|
_, err := txn.Get(keyBytes)
|
|
if err != nil {
|
|
if errors.Is(err, badger.ErrKeyNotFound) {
|
|
wasSet = true
|
|
return wrapErr(txn.Set(keyBytes, value))
|
|
}
|
|
}
|
|
return wrapErr(err)
|
|
})
|
|
return wasSet, wrapErr(err)
|
|
}
|
|
|
|
func (k *keyValueStorage) SetIfNotExistWithExpiration(ctx context.Context, key string, value []byte, expiration time.Time) (bool, error) {
|
|
var wasSet bool
|
|
err := k.db.Update(func(txn *badger.Txn) error {
|
|
keyBytes := []byte(k.prefix + key)
|
|
_, err := txn.Get(keyBytes)
|
|
if err != nil {
|
|
if errors.Is(err, badger.ErrKeyNotFound) {
|
|
wasSet = true
|
|
entry := badger.NewEntry(keyBytes, value)
|
|
entry.ExpiresAt = uint64(expiration.Unix())
|
|
return wrapErr(txn.SetEntry(entry))
|
|
}
|
|
}
|
|
return wrapErr(err)
|
|
})
|
|
return wasSet, err
|
|
}
|
|
|
|
func (k *keyValueStorage) ExpireAt(ctx context.Context, key string, expiration time.Time) error {
|
|
return k.db.Update(func(txn *badger.Txn) error {
|
|
keyBytes := []byte(k.prefix + key)
|
|
item, err := txn.Get(keyBytes)
|
|
if err != nil {
|
|
return wrapErr(err)
|
|
}
|
|
|
|
copyOfBytes := make([]byte, item.ValueSize())
|
|
_, err = item.ValueCopy(copyOfBytes)
|
|
if err != nil {
|
|
return wrapErr(err)
|
|
}
|
|
|
|
entry := badger.NewEntry(keyBytes, copyOfBytes)
|
|
entry.ExpiresAt = uint64(expiration.Unix())
|
|
return wrapErr(txn.SetEntry(entry))
|
|
})
|
|
}
|
|
|
|
func (k *keyValueStorage) Remove(ctx context.Context, key string) error {
|
|
return k.db.Update(func(txn *badger.Txn) error {
|
|
keyBytes := []byte(k.prefix + key)
|
|
return wrapErr(txn.Delete(keyBytes))
|
|
})
|
|
}
|
|
|
|
func (k *keyValueStorage) HasCapability(capability kv.Capability) bool {
|
|
switch capability {
|
|
default:
|
|
return false
|
|
}
|
|
}
|