86 lines
1.5 KiB
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()
|
||
|
|
}
|