pelican/internal/db/database.go

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 dbWriteMode int

const (
	WriteModeCreate dbWriteMode = iota
	WriteModeUpdate
)

type BoltItem interface {
	SetID(int) error
	GetID() 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 the Bolt item; %w", err)
		}

		return nil
	})
	if err != nil {
		return nil, fmt.Errorf("error while loading the Bolt items from the database; %w", err)
	}

	return output, nil
}

// Write saves one or more Bolt items to the status bucket.
func Write(database *bolt.DB, writeMode dbWriteMode, bucketName string, items []BoltItem) ([]int, error) {
	if len(items) == 0 {
		return nil, itemListEmptyError{}
	}

	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 writeMode == WriteModeCreate && item.GetID() < 1 {
				var id uint64
				if id, err = bucket.NextSequence(); err != nil {
					return fmt.Errorf("unable to generate ID, %w", err)
				}

				if err = item.SetID(int(id)); err != nil {
					return fmt.Errorf("unable to set the generated ID to the Bolt item; %w", err)
				}
			}

			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.GetID())), buf.Bytes()); err != nil {
				return fmt.Errorf("unable to add the Bolt Item to the %s bucket, %w", bucketName, err)
			}

			ids[ind] = item.GetID()
		}

		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
}