pelican/db.go

package main

import (
	"bytes"
	"encoding/gob"
	"fmt"
	"os"
	"path/filepath"
	"runtime"
	"strconv"
	"time"

	bolt "go.etcd.io/bbolt"
)

const (
	statusBucket string = "status"
	cardBucket   string = "card"
)

func mkDataDir(dir string) error {
	return os.MkdirAll(dir, 0700)
}

// TODO: Needs unit testing and documentation
func dbPath(path string) (string, error) {
	if len(path) > 0 {
		filepath.Dir(path)
		return path, mkDataDir(filepath.Dir(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")
	}

	if err := os.MkdirAll(dataDir, 0700); err != nil {
		return "", fmt.Errorf("unable to make directory %s, %w", dataDir, err)
	}

	path = filepath.Join(dataDir, dbFilename)

	return path, mkDataDir(dataDir)

}

// openDatabase opens the database, at a given path, for reading and writing. If the file does not exist it will be created.
// For linux, the default location of the database file is $XDG_DATA_HOME/pelican/pelican.db. If the XDG_DATA_HOME environment
// variable is not set then it will default to $HOME/.local/share/pelican/pelican.db.
// For all other operating systems the default location is $HOME/.pelican/pelican.db.
func openDatabase(path string) (*bolt.DB, error) {
	opts := bolt.Options{
		Timeout: 1 * time.Second,
	}

	db, err := bolt.Open(path, 0600, &opts)
	if err != nil {
		return nil, fmt.Errorf("unable to open database at %s, %w", path, err)
	}

	return db, nil
}

// ensureBuckets ensures that the required buckets are created in the database.
func ensureBuckets(db *bolt.DB) error {
	buckets := []string{statusBucket, cardBucket}

	return 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
	})
}

// saveStatuses saves one or more statuses to the status bucket.
func saveStatuses(db *bolt.DB, statuses []Status) error {
	if len(statuses) == 0 {
		return nil
	}

	bucket := []byte(statusBucket)

	return db.Update(func(tx *bolt.Tx) error {
		b := tx.Bucket(bucket)

		if b == nil {
			return fmt.Errorf("bucket %s does not exist", bucket)
		}

		for _, v := range statuses {
			var err error

			if v.Id == 0 {
				var id uint64
				if id, err = b.NextSequence(); err != nil {
					return fmt.Errorf("unable to generate ID, %w", err)
				}
				v.Id = int(id)
			}

			buf := new(bytes.Buffer)
			encoder := gob.NewEncoder(buf)
			if err = encoder.Encode(v); err != nil {
				return fmt.Errorf("unable to encode data, %w", err)
			}

			if err = b.Put([]byte(strconv.Itoa(v.Id)), buf.Bytes()); err != nil {
				return fmt.Errorf("unable to add the status to the bucket, %w", err)
			}
		}

		return nil
	})
}

// loadAllStatuses retrieves all the statuses from the status bucket.
func loadAllStatuses(db *bolt.DB) ([]Status, error) {
	var statuses []Status
	bucket := []byte(statusBucket)

	err := db.View(func(tx *bolt.Tx) error {
		b := tx.Bucket(bucket)

		if b == nil {
			return fmt.Errorf("bucket %s does not exist", bucket)
		}

		if err := b.ForEach(func(_, v []byte) error {
			var s Status
			buf := bytes.NewBuffer(v)
			decoder := gob.NewDecoder(buf)
			if err := decoder.Decode(&s); err != nil {
				return fmt.Errorf("unable to decode data, %w", err)
			}

			statuses = append(statuses, s)

			return nil
		}); err != nil {
			return err
		}
		return nil
	})

	return statuses, err
}

// saveCard writes a card to the card bucket in BoltDB.
func saveCard(db *bolt.DB, card Card) (int, error) {
	bucket := []byte(cardBucket)

	err := db.Update(func(tx *bolt.Tx) error {
		var err error
		b := tx.Bucket(bucket)

		if b == nil {
			return fmt.Errorf("bucket %s does not exist", bucket)
		}

		if card.Id == 0 {
			var id uint64
			if id, err = b.NextSequence(); err != nil {
				return fmt.Errorf("unable to generate an ID for the card, %w", err)
			}
			card.Id = int(id)
		}

		buf := new(bytes.Buffer)
		encoder := gob.NewEncoder(buf)
		if err = encoder.Encode(card); err != nil {
			return fmt.Errorf("unable to encode data, %w", err)
		}

		if err = b.Put([]byte(strconv.Itoa(card.Id)), buf.Bytes()); err != nil {
			return fmt.Errorf("unable to write the card to the bucket, %w", err)
		}

		return nil
	})

	return card.Id, err
}

// loadCard retrieves a card from the cards bucket in BoltDB.
func loadCard(db *bolt.DB, id int) (Card, error) {
	var card Card
	bucket := []byte(cardBucket)

	err := db.View(func(tx *bolt.Tx)error {
		b := tx.Bucket(bucket)

		if b == nil {
			return fmt.Errorf("bucket %s does not exist", bucket)
		}

		buf := bytes.NewBuffer(b.Get([]byte(strconv.Itoa(id))))

		decoder := gob.NewDecoder(buf)
		if err := decoder.Decode(&card); err != nil {
			return fmt.Errorf("unable to decode card data, %w", err)
		}

		return nil
	})

	return card, err
}