doomer/db.go

86 lines
1.5 KiB
Go

package main
import (
"time"
badger "github.com/dgraph-io/badger/v3"
)
type Db struct {
db *badger.DB
}
func NewDatabase() (*Db, error) {
db, err := badger.Open(badger.DefaultOptions("./database"))
if err != nil {
return nil, err
}
return &Db{
db: db,
}, nil
}
type PrefixStore struct {
db *badger.DB
prefix ItemType
}
type ItemType string
const (
GenericItem ItemType = ""
ImageItem ItemType = "images:"
FileIDItem ItemType = "fileid:"
)
func (d *Db) Prefix(prefix ItemType) PrefixStore {
return PrefixStore{
db: d.db,
prefix: prefix,
}
}
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) 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 image []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 {
image = append(image, val...)
return nil
})
})
if err != nil {
return nil, err
}
return image, nil
}
func (d *Db) Close() {
d.db.Close()
}