package db import ( "bytes" "encoding/gob" "fmt" "os" "path/filepath" "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 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 } // 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 }