61 lines
1.1 KiB
Go
61 lines
1.1 KiB
Go
|
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)))
|
||
|
|
})
|
||
|
|
}
|