stonebot/database/database.go

70 lines
1.3 KiB
Go

package database
import (
"errors"
"os"
"strconv"
"go.etcd.io/bbolt"
)
var (
stonelandBucket = []byte("stoneland")
stonesKey = []byte("stones")
)
type Database struct {
db *bbolt.DB
}
func New(dbPath string) (*Database, error) {
db, err := bbolt.Open(dbPath, os.ModePerm, bbolt.DefaultOptions)
if err != nil {
return nil, err
}
return &Database{
db: db,
}, db.Update(func(tx *bbolt.Tx) error {
_, err := tx.CreateBucketIfNotExists(stonelandBucket)
return err
})
}
func (db *Database) Close() error {
return db.db.Close()
}
func (db *Database) Rock(num int) error {
return db.db.Update(func(tx *bbolt.Tx) error {
bucket := tx.Bucket(stonelandBucket)
if bucket == nil {
return errors.New("bucket for stones was not created")
}
stones, err := db.rocks(tx)
if err != nil {
return err
}
stones += num
return bucket.Put(stonesKey, []byte(strconv.Itoa(stones)))
})
}
func (db *Database) Rocks() (count int, err error) {
return count, db.db.View(func(tx *bbolt.Tx) error {
count, err = db.rocks(tx)
return err
})
}
func (db *Database) rocks(tx *bbolt.Tx) (int, error) {
bucket := tx.Bucket(stonelandBucket)
if bucket == nil {
return -1, errors.New("bucket for stones was not created")
}
stones := bucket.Get(stonesKey)
if stones == nil {
return 0, nil
}
return strconv.Atoi(string(stones))
}