canopeas/database/database.go

61 lines
1.1 KiB
Go
Raw Permalink Normal View History

package database
import (
"strconv"
"go.etcd.io/bbolt"
)
var (
firedBucket = []byte("fired")
buckets = [][]byte{
firedBucket,
}
)
type Database struct {
db *bbolt.DB
}
func Load(dbPath string) (*Database, error) {
db, err := bbolt.Open(dbPath, 0600, bbolt.DefaultOptions)
if err != nil {
return nil, err
}
for _, b := range buckets {
if err := db.Update(func(tx *bbolt.Tx) error {
_, err := tx.CreateBucketIfNotExists(b)
return err
}); err != nil {
return nil, err
}
}
return &Database{
db: db,
}, nil
}
func (db *Database) CheckPing(roleID string) int {
roleIDByte := []byte(roleID)
var idx int
_ = db.db.View(func(tx *bbolt.Tx) error {
num := tx.Bucket(firedBucket).Get(roleIDByte)
if num != nil {
if i, err := strconv.Atoi(string(num)); err == nil {
idx = i
}
}
return nil
})
return idx
}
func (db *Database) IncrementPing(roleID string) error {
roleIDByte := []byte(roleID)
return db.db.Update(func(tx *bbolt.Tx) error {
idx := db.CheckPing(roleID)
return tx.Bucket(firedBucket).Put(roleIDByte, []byte(strconv.Itoa(idx+1)))
})
}