310 lines
7.3 KiB
Go
310 lines
7.3 KiB
Go
|
package db
|
||
|
|
||
|
import (
|
||
|
"bytes"
|
||
|
"encoding/gob"
|
||
|
"fmt"
|
||
|
"os"
|
||
|
"path/filepath"
|
||
|
"runtime"
|
||
|
"strconv"
|
||
|
"time"
|
||
|
|
||
|
bolt "go.etcd.io/bbolt"
|
||
|
)
|
||
|
|
||
|
const (
|
||
|
StatusBucket string = "status"
|
||
|
CardBucket string = "card"
|
||
|
)
|
||
|
|
||
|
type BoltItem interface {
|
||
|
UpdateId(int)
|
||
|
Id() int
|
||
|
}
|
||
|
|
||
|
// OpenDatabase opens the database, at a given path, for reading and writing.
|
||
|
// If the file does not exist it will be created.
|
||
|
func OpenDatabase(path string) (*bolt.DB, error) {
|
||
|
var err error
|
||
|
|
||
|
path = dbPath(path)
|
||
|
|
||
|
if err = mkDataDir(path); err != nil {
|
||
|
return nil, fmt.Errorf("unable to make the data directory, %w", err)
|
||
|
}
|
||
|
|
||
|
opts := bolt.Options{
|
||
|
Timeout: 1 * time.Second,
|
||
|
}
|
||
|
|
||
|
var database *bolt.DB
|
||
|
|
||
|
if database, err = bolt.Open(path, 0o600, &opts); err != nil {
|
||
|
return nil, fmt.Errorf("unable to open database at %s, %w", path, err)
|
||
|
}
|
||
|
|
||
|
if err = ensureBuckets(database); err != nil {
|
||
|
return nil, fmt.Errorf("unable to ensure the required buckets are in the database, %w", err)
|
||
|
}
|
||
|
|
||
|
return database, nil
|
||
|
}
|
||
|
|
||
|
// Read retrieves a Bolt item from a specified bucket and returns the data in bytes.
|
||
|
func Read(db *bolt.DB, bucketName string, itemID int) ([]byte, error) {
|
||
|
bucketNameBytes := []byte(bucketName)
|
||
|
|
||
|
var data []byte
|
||
|
|
||
|
if err := db.View(func(tx *bolt.Tx) error {
|
||
|
bucket := tx.Bucket(bucketNameBytes)
|
||
|
|
||
|
if bucket == nil {
|
||
|
return bucketNotExistError{bucket: string(bucketNameBytes)}
|
||
|
}
|
||
|
|
||
|
data = bucket.Get([]byte(strconv.Itoa(itemID)))
|
||
|
|
||
|
return nil
|
||
|
}); err != nil {
|
||
|
return nil, fmt.Errorf("error while reading the Bolt item from the database, %w", err)
|
||
|
}
|
||
|
|
||
|
return data, nil
|
||
|
}
|
||
|
|
||
|
// ReadMany reads one or more Bolt items from the specified bucket.
|
||
|
func ReadMany(db *bolt.DB, bucketName string, ids []int) ([][]byte, error) {
|
||
|
bucketNameBytes := []byte(bucketName)
|
||
|
|
||
|
var output [][]byte
|
||
|
|
||
|
err := db.View(func(tx *bolt.Tx) error {
|
||
|
bucket := tx.Bucket(bucketNameBytes)
|
||
|
|
||
|
if bucket == nil {
|
||
|
return bucketNotExistError{bucket: bucketName}
|
||
|
}
|
||
|
|
||
|
for _, v := range ids {
|
||
|
data := bucket.Get([]byte(strconv.Itoa(v)))
|
||
|
output = append(output, data)
|
||
|
}
|
||
|
|
||
|
return nil
|
||
|
})
|
||
|
if err != nil {
|
||
|
return nil, fmt.Errorf("error while retrieving the data from the database, %w", err)
|
||
|
}
|
||
|
|
||
|
return output, nil
|
||
|
}
|
||
|
|
||
|
// ReadAll retrieves all the Bolt Items from the specified bucket.
|
||
|
func ReadAll(db *bolt.DB, bucketName string) ([][]byte, error) {
|
||
|
bucketNameBytes := []byte(bucketName)
|
||
|
|
||
|
var output [][]byte
|
||
|
|
||
|
err := db.View(func(tx *bolt.Tx) error {
|
||
|
bucket := tx.Bucket(bucketNameBytes)
|
||
|
|
||
|
if bucket == nil {
|
||
|
return bucketNotExistError{bucket: bucketName}
|
||
|
}
|
||
|
|
||
|
if err := bucket.ForEach(func(_, v []byte) error {
|
||
|
output = append(output, v)
|
||
|
|
||
|
return nil
|
||
|
}); err != nil {
|
||
|
return fmt.Errorf("unable to load status, %w", err)
|
||
|
}
|
||
|
|
||
|
return nil
|
||
|
})
|
||
|
if err != nil {
|
||
|
return nil, fmt.Errorf("error while loading statuses from the database, %w", err)
|
||
|
}
|
||
|
|
||
|
return output, nil
|
||
|
}
|
||
|
|
||
|
// Write creates or updates a Bolt item to a specified bucket.
|
||
|
func Write(db *bolt.DB, bucketName string, item BoltItem) (int, error) {
|
||
|
bucketNameBytes := []byte(bucketName)
|
||
|
|
||
|
err := db.Update(func(tx *bolt.Tx) error {
|
||
|
var err error
|
||
|
bucket := tx.Bucket(bucketNameBytes)
|
||
|
|
||
|
if bucket == nil {
|
||
|
return bucketNotExistError{bucket: string(bucketNameBytes)}
|
||
|
}
|
||
|
|
||
|
if item.Id() < 1 {
|
||
|
var id uint64
|
||
|
if id, err = bucket.NextSequence(); err != nil {
|
||
|
return fmt.Errorf("unable to generate an ID for the card, %w", err)
|
||
|
}
|
||
|
item.UpdateId(int(id))
|
||
|
}
|
||
|
|
||
|
buf := new(bytes.Buffer)
|
||
|
encoder := gob.NewEncoder(buf)
|
||
|
if err = encoder.Encode(item); err != nil {
|
||
|
return fmt.Errorf("unable to encode data, %w", err)
|
||
|
}
|
||
|
|
||
|
if err = bucket.Put([]byte(strconv.Itoa(item.Id())), buf.Bytes()); err != nil {
|
||
|
return fmt.Errorf("unable to write the card to the bucket, %w", err)
|
||
|
}
|
||
|
|
||
|
return nil
|
||
|
})
|
||
|
if err != nil {
|
||
|
return 0, fmt.Errorf("error while saving the card to the database, %w", err)
|
||
|
}
|
||
|
|
||
|
return item.Id(), nil
|
||
|
}
|
||
|
|
||
|
// WriteMany saves one or more Bolt items to the status bucket.
|
||
|
func WriteMany(database *bolt.DB, bucketName string, items []BoltItem) ([]int, error) {
|
||
|
if len(items) == 0 {
|
||
|
return []int{}, nil
|
||
|
}
|
||
|
|
||
|
ids := make([]int, len(items))
|
||
|
|
||
|
bucketNameBytes := []byte(bucketName)
|
||
|
|
||
|
err := database.Update(func(tx *bolt.Tx) error {
|
||
|
bucket := tx.Bucket(bucketNameBytes)
|
||
|
|
||
|
if bucket == nil {
|
||
|
return bucketNotExistError{bucket: string(bucketNameBytes)}
|
||
|
}
|
||
|
|
||
|
for ind, item := range items {
|
||
|
var err error
|
||
|
|
||
|
if item.Id() < 1 {
|
||
|
var id uint64
|
||
|
if id, err = bucket.NextSequence(); err != nil {
|
||
|
return fmt.Errorf("unable to generate ID, %w", err)
|
||
|
}
|
||
|
item.UpdateId(int(id))
|
||
|
}
|
||
|
|
||
|
buf := new(bytes.Buffer)
|
||
|
encoder := gob.NewEncoder(buf)
|
||
|
if err = encoder.Encode(item); err != nil {
|
||
|
return fmt.Errorf("unable to encode data, %w", err)
|
||
|
}
|
||
|
|
||
|
if err = bucket.Put([]byte(strconv.Itoa(item.Id())), buf.Bytes()); err != nil {
|
||
|
return fmt.Errorf("unable to add the Bolt Item to the %s bucket, %w", bucketName, err)
|
||
|
}
|
||
|
|
||
|
ids[ind] = item.Id()
|
||
|
}
|
||
|
|
||
|
return nil
|
||
|
})
|
||
|
if err != nil {
|
||
|
return nil, fmt.Errorf("error while saving the Bolt Items to the database, %w", err)
|
||
|
}
|
||
|
|
||
|
return ids, nil
|
||
|
}
|
||
|
|
||
|
// Delete deletes a Bolt item from a specified bucket.
|
||
|
func Delete(db *bolt.DB, bucketName string, itemID int) error {
|
||
|
bucketNameBytes := []byte(bucketName)
|
||
|
|
||
|
if err := db.Update(func(tx *bolt.Tx) error {
|
||
|
bucket := tx.Bucket(bucketNameBytes)
|
||
|
|
||
|
if bucket == nil {
|
||
|
return bucketNotExistError{bucket: bucketName}
|
||
|
}
|
||
|
|
||
|
if err := bucket.Delete([]byte(strconv.Itoa(itemID))); err != nil {
|
||
|
return fmt.Errorf("an error occurred when deleting Bolt item '%d', %w", itemID, err)
|
||
|
}
|
||
|
|
||
|
return nil
|
||
|
}); err != nil {
|
||
|
return fmt.Errorf("error deleting data from the '%s' bucket, %w", bucketName, err)
|
||
|
}
|
||
|
|
||
|
return nil
|
||
|
}
|
||
|
|
||
|
// dbPath returns the path to the database file. If a path is given then that is returned. Otherwise the default path is returned.
|
||
|
// For linux, the default location of the database file is $XDG_DATA_HOME/canal/canal.db. If the XDG_DATA_HOME environment
|
||
|
// variable is not set then it will default to $HOME/.local/share/canal/canal.db. For all other operating systems the default
|
||
|
// location is $HOME/.canal/canal.db.
|
||
|
func dbPath(path string) string {
|
||
|
if len(path) > 0 {
|
||
|
filepath.Dir(path)
|
||
|
|
||
|
return path
|
||
|
}
|
||
|
|
||
|
dbFilename := "pelican.db"
|
||
|
|
||
|
var dataDir string
|
||
|
|
||
|
goos := runtime.GOOS
|
||
|
switch goos {
|
||
|
case "linux":
|
||
|
dataHome := os.Getenv("XDG_DATA_HOME")
|
||
|
|
||
|
if len(dataHome) == 0 {
|
||
|
dataHome = filepath.Join(os.Getenv("HOME"), ".local", "share")
|
||
|
}
|
||
|
|
||
|
dataDir = filepath.Join(dataHome, "pelican")
|
||
|
default:
|
||
|
dataDir = filepath.Join(os.Getenv("HOME"), ".pelican")
|
||
|
}
|
||
|
|
||
|
path = filepath.Join(dataDir, dbFilename)
|
||
|
|
||
|
return path
|
||
|
}
|
||
|
|
||
|
// mkDataDir creates the data directory of a given path to the database.
|
||
|
func mkDataDir(path string) error {
|
||
|
dir := filepath.Dir(path)
|
||
|
|
||
|
if err := os.MkdirAll(dir, 0o700); err != nil {
|
||
|
return fmt.Errorf("error while making directory %s, %w", dir, err)
|
||
|
}
|
||
|
|
||
|
return nil
|
||
|
}
|
||
|
|
||
|
// ensureBuckets ensures that the required buckets are created in the database.
|
||
|
func ensureBuckets(db *bolt.DB) error {
|
||
|
buckets := []string{StatusBucket, CardBucket}
|
||
|
|
||
|
err := db.Update(func(tx *bolt.Tx) error {
|
||
|
for _, v := range buckets {
|
||
|
if _, err := tx.CreateBucketIfNotExists([]byte(v)); err != nil {
|
||
|
return fmt.Errorf("unable to ensure that %s bucket is created in the database, %w", v, err)
|
||
|
}
|
||
|
}
|
||
|
|
||
|
return nil
|
||
|
})
|
||
|
if err != nil {
|
||
|
return fmt.Errorf("error while ensuring buckets exist in the database, %w", err)
|
||
|
}
|
||
|
|
||
|
return nil
|
||
|
}
|