122 lines
2.2 KiB
Go
122 lines
2.2 KiB
Go
package db
|
|
|
|
import (
|
|
"time"
|
|
|
|
badger "github.com/dgraph-io/badger/v3"
|
|
)
|
|
|
|
type Db struct {
|
|
db *badger.DB
|
|
}
|
|
|
|
var database Db
|
|
|
|
func NewDatabase() (*Db, error) {
|
|
db, err := badger.Open(badger.DefaultOptions("./database"))
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
database = Db{
|
|
db: db,
|
|
}
|
|
return &database, nil
|
|
}
|
|
|
|
func DB() *Db {
|
|
return &database
|
|
}
|
|
|
|
type PrefixStore struct {
|
|
db *badger.DB
|
|
prefix ItemType
|
|
}
|
|
|
|
type ItemType string
|
|
|
|
const (
|
|
GenericItem ItemType = ""
|
|
ImageItem ItemType = "images:"
|
|
FileIDItem ItemType = "fileid:"
|
|
ConfigItem ItemType = "config:"
|
|
)
|
|
|
|
func (d *Db) Prefix(prefix ItemType) PrefixStore {
|
|
return PrefixStore{
|
|
db: d.db,
|
|
prefix: prefix,
|
|
}
|
|
}
|
|
|
|
func (d *PrefixStore) Suffix(suffix string) PrefixStore {
|
|
return PrefixStore{
|
|
db: d.db,
|
|
prefix: d.prefix + ItemType(suffix),
|
|
}
|
|
}
|
|
|
|
func (d *PrefixStore) StoreItem(key string, input []byte) error {
|
|
fullKey := string(d.prefix) + key
|
|
return d.db.Update(func(txn *badger.Txn) error {
|
|
e := badger.NewEntry([]byte(fullKey), input).WithTTL(time.Hour * 60 * 24 * 10)
|
|
return txn.SetEntry(e)
|
|
})
|
|
}
|
|
|
|
func (d *PrefixStore) StoreItemStr(key string, input string) error {
|
|
return d.StoreItem(key, []byte(input))
|
|
}
|
|
|
|
func (d *PrefixStore) ItemExists(key string) bool {
|
|
exists := false
|
|
fullKey := string(d.prefix) + key
|
|
d.db.View(func(txn *badger.Txn) error {
|
|
_, err := txn.Get([]byte(fullKey))
|
|
if err == nil {
|
|
exists = true
|
|
}
|
|
return err
|
|
})
|
|
return exists
|
|
}
|
|
|
|
func (d *PrefixStore) GetItem(key string) ([]byte, error) {
|
|
var item []byte
|
|
fullKey := string(d.prefix) + key
|
|
err := d.db.View(func(txn *badger.Txn) error {
|
|
result, err := txn.Get([]byte(fullKey))
|
|
if err != nil {
|
|
return err
|
|
}
|
|
return result.Value(func(val []byte) error {
|
|
item = append(item, val...)
|
|
return nil
|
|
})
|
|
})
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return item, nil
|
|
}
|
|
|
|
func (d *PrefixStore) GetItemStr(key string) (string, error) {
|
|
item, err := d.GetItem(key)
|
|
return string(item), err
|
|
}
|
|
|
|
func (d *PrefixStore) GetItemOr(key string, def []byte) []byte {
|
|
item, err := d.GetItem(key)
|
|
if err != nil {
|
|
return def
|
|
}
|
|
return item
|
|
}
|
|
|
|
func (d *PrefixStore) GetItemOrStr(key string, def string) string {
|
|
return string(d.GetItemOr(key, []byte(def)))
|
|
}
|
|
|
|
func (d *Db) Close() {
|
|
d.db.Close()
|
|
}
|