rushlink/db.go

package rushlink

import (
	"fmt"
	"log"
	"time"

	"github.com/pkg/errors"
	bolt "go.etcd.io/bbolt"
)

var DB *bolt.DB

// The current database version
//
// If we alter the database format, we bump this number and write a new
// database migration in migrateDatabase().
const CURRENT_MIGRATE_VERSION = 1

// Bucket storing everything that is not a bulk value. This includes stuff like
// the database version, secret site-wide keys.
const BUCKET_CONF = "conf"

// The main bucket for paste values and URL redirects
const BUCKET_PASTES = "pastes"

// This value stores the current migration version. If this value is less than
// CURRENT_MIGRATE_VERSION, the database has to be migrated.
const KEY_MIGRATE_VERSION = "migrate_version"

// Open the bolt database
func Open(path string) error {
	if path == "" {
		return errors.New("database not set")
	}

	var err error
	DB, err = bolt.Open(path, 0666, &bolt.Options{Timeout: 1 * time.Second})
	if err != nil {
		return errors.Wrapf(err, "failed to open database at '%v'", path)
	}
	return DB.Update(migrateDatabase)
}

// Close the bolt database
func Close() error {
	if DB == nil {
		panic("no open database")
	}
	return DB.Close()
}

// Initialize and migrate the database to the current version
func migrateDatabase(tx *bolt.Tx) error {
	dbVersion, err := dbVersion(tx)
	if err != nil {
		return err
	}

	// Migrate the database to version 1
	if dbVersion < 1 {
		log.Println("migrating database to version 1")

		// Create conf bucket
		_, err := tx.CreateBucket([]byte(BUCKET_CONF))
		if err != nil {
			return err
		}

		// Create paste bucket
		_, err = tx.CreateBucket([]byte(BUCKET_PASTES))
		if err != nil {
			return err
		}

		// Update the version number
		if err := setDBVersion(tx, 1); err != nil {
			return err
		}
	}

	return nil
}

// Get the current migrate version from the database
func dbVersion(tx *bolt.Tx) (int, error) {
	conf := tx.Bucket([]byte(BUCKET_CONF))
	if conf == nil {
		return 0, nil
	}
	dbVersionBytes := conf.Get([]byte(KEY_MIGRATE_VERSION))
	if dbVersionBytes == nil {
		return 0, nil
	}

	// Version was already stored
	var dbVersion int
	if err := Unmarshal(dbVersionBytes, &dbVersion); err != nil {
		return 0, err
	}
	if dbVersion == 0 {
		return 0, fmt.Errorf("database version is invalid (%v)", dbVersion)
	}
	if dbVersion > CURRENT_MIGRATE_VERSION {
		return 0, fmt.Errorf("database version is too recent (%v > %v)", dbVersion, CURRENT_MIGRATE_VERSION)
	}
	return dbVersion, nil
}

// Update the current migrate version in the database
func setDBVersion(tx *bolt.Tx, version int) error {
	conf, err := tx.CreateBucketIfNotExists([]byte(BUCKET_CONF))
	if err != nil {
		return err
	}

	versionBytes, err := Marshal(version)
	if err != nil {
		return err
	}
	return conf.Put([]byte(KEY_MIGRATE_VERSION), versionBytes)
}